diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..8b880bff56 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -17,6 +17,7 @@ TrainEntrypointOptions, ) from .metrics import ( + MetricAccumulator, TrainingMetricAccumulator, ) from .schedule import ( @@ -50,6 +51,7 @@ "CheckpointStore", "DisplayInterval", "LearningCurveWriter", + "MetricAccumulator", "RankContext", "ShardingPolicy", "StepSchedule", diff --git a/deepmd/dpmodel/train/metrics.py b/deepmd/dpmodel/train/metrics.py index 3ea0f405b8..2996e053d6 100644 --- a/deepmd/dpmodel/train/metrics.py +++ b/deepmd/dpmodel/train/metrics.py @@ -1,23 +1,77 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Backend-independent accumulation of training display metrics.""" +"""Backend-independent accumulation of training display metrics. -from collections.abc import ( - Mapping, - Sequence, +Two weighting policies share one accumulator: + +* Training steps use unit weight, so the reported value is the arithmetic + mean of the detached per-step metrics. +* Validation batches use an atom weight (real atoms in the batch), so the + reported value is the atom-weighted mean of per-atom metrics. + +Optional labels emit host NaN through ``display_if_exist``. Those observations +are excluded from that metric's sum and weight so a single unlabeled batch +cannot poison an otherwise finite interval. Present device metrics are never +host-synced on ``add``; distributed runs keep local sums and weights on device +and reduce them only at a display boundary. +""" + +from __future__ import ( + annotations, +) + +from math import ( + isnan, ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from collections.abc import ( + Mapping, + Sequence, + ) + -class TrainingMetricAccumulator: - """Average per-step metrics independently for each training task. +def _is_host_nan(value: Any) -> bool: + """Return whether a host scalar is NaN without touching device memory. - Metric names are declared before training so an unsampled task has the - same columns as a sampled task. Each update has equal weight, matching - the arithmetic mean of the reported per-step metrics rather than a - pooled error over atoms. Scalar arrays stay on their original device - until :meth:`average` converts them to Python floats for display. + Optional-label gaps arrive as host ``float("nan")`` from + ``display_if_exist``. Present CUDA / accelerator metrics must stay on + device through :meth:`MetricAccumulator.add`; coercing them with + ``float()`` would D2H-sync on every training step. Device values are + therefore never inspected here. Host conversion happens only in + :meth:`average` and in display-boundary all-reduce packing. + """ + if value is None: + return True + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return isnan(float(value)) + # 0-d NumPy arrays and NumPy scalars are already on the host. + if type(value).__module__.startswith("numpy"): + ndim = getattr(value, "ndim", 0) + if ndim == 0: + try: + return bool(isnan(float(value))) + except (TypeError, ValueError): + return False + return False + + +class MetricAccumulator: + """Accumulate weighted sums and counts of detached display metrics. + + Metric names are declared before training so an unsampled task keeps the + same columns as a sampled task. Each call to :meth:`add` is one + observation: training uses ``weight=1``, validation passes the number of + real atoms. A host-NaN metric is skipped for that observation only. + + Scalar arrays stay on their original device until :meth:`average` + converts them to Python floats for display. Out-of-place addition + preserves recorded values when a backend reuses an output buffer. Parameters ---------- @@ -26,44 +80,123 @@ class TrainingMetricAccumulator: """ def __init__(self, metric_names: Mapping[str, Sequence[str]]) -> None: - self._totals: dict[str, dict[str, Any]] = { - task: dict.fromkeys(sorted(names), 0.0) - for task, names in metric_names.items() + self._names: dict[str, tuple[str, ...]] = { + task: tuple(sorted(names)) for task, names in metric_names.items() + } + self._sums: dict[str, dict[str, Any]] = { + task: dict.fromkeys(names, 0.0) for task, names in self._names.items() + } + self._weights: dict[str, dict[str, Any]] = { + task: dict.fromkeys(names, 0.0) for task, names in self._names.items() } - self._counts = dict.fromkeys(metric_names, 0) + self._observations = dict.fromkeys(self._names, 0) - def add(self, task_key: str, metrics: Mapping[str, Any]) -> None: - """Record one optimizer step's detached scalar metrics. + @property + def task_keys(self) -> tuple[str, ...]: + """Return the declared task keys in insertion order.""" + return tuple(self._names) + + def metric_names(self, task_key: str) -> tuple[str, ...]: + """Return the declared metric names for one task.""" + return self._names[task_key] + + def add( + self, + task_key: str, + metrics: Mapping[str, Any], + *, + weight: float = 1.0, + ) -> None: + """Record one observation of detached scalar metrics. Backends detach metrics from their differentiation graph before - calling this method. Missing metrics propagate NaN rather than being - treated as zero. Out-of-place addition preserves the recorded value - when a backend reuses an output buffer on its next step. + calling this method. Missing or host-NaN metrics are omitted from that + metric's weighted average rather than treated as zero. A non-positive + weight records the observation for task sampling but contributes to + no metric totals. """ - totals = self._totals[task_key] - unknown = metrics.keys() - totals.keys() + names = self._names[task_key] + unknown = metrics.keys() - set(names) if unknown: raise ValueError( f"Undeclared display metrics for task {task_key!r}: {sorted(unknown)}" ) - for name, total in totals.items(): - totals[name] = total + metrics.get(name, float("nan")) - self._counts[task_key] += 1 + sums = self._sums[task_key] + weights = self._weights[task_key] + for name in names: + value = metrics.get(name, float("nan")) + if weight <= 0.0 or _is_host_nan(value): + continue + # Out-of-place update: backends may overwrite the tensor buffer + # that produced ``value`` on the next step. + sums[name] = sums[name] + value * weight + weights[name] = weights[name] + weight + self._observations[task_key] += 1 def count(self, task_key: str) -> int: - """Return the number of recorded optimizer steps for one task.""" - return self._counts[task_key] + """Return the number of recorded observations for one task. + + This is the sampling count used to detect tasks that never ran in a + display interval. It is independent of per-metric weights. + """ + return self._observations[task_key] def average(self, task_key: str) -> dict[str, float]: - """Return interval averages, or NaN for every unsampled metric.""" - totals = self._totals[task_key] - count = self._counts[task_key] - if count == 0: - return dict.fromkeys(totals, float("nan")) - return {name: float(total / count) for name, total in totals.items()} + """Return weighted averages, or NaN for every metric with zero weight.""" + sums = self._sums[task_key] + weights = self._weights[task_key] + result: dict[str, float] = {} + for name in self._names[task_key]: + weight = float(weights[name]) + if weight == 0.0 or isnan(weight): + result[name] = float("nan") + else: + result[name] = float(sums[name] / weights[name]) + return result def reset(self) -> None: """Clear interval values and counts while preserving metric names.""" - for task, totals in self._totals.items(): - self._totals[task] = dict.fromkeys(totals, 0.0) - self._counts[task] = 0 + for task, names in self._names.items(): + self._sums[task] = dict.fromkeys(names, 0.0) + self._weights[task] = dict.fromkeys(names, 0.0) + self._observations[task] = 0 + + def flat_state(self) -> list[Any]: + """Pack sums, weights, and observation counts for a collective reduce. + + Layout is stable across ranks that share the same constructor + arguments: for each task, all metric sums, then all metric weights, + then the observation count. + """ + values: list[Any] = [] + for task, names in self._names.items(): + sums = self._sums[task] + weights = self._weights[task] + values.extend(sums[name] for name in names) + values.extend(weights[name] for name in names) + values.append(float(self._observations[task])) + return values + + def load_flat_state(self, values: Sequence[Any]) -> None: + """Replace local state with a previously packed, possibly reduced vector.""" + expected = sum(2 * len(names) + 1 for names in self._names.values()) + if len(values) != expected: + raise ValueError( + f"MetricAccumulator state length {len(values)} != expected {expected}" + ) + index = 0 + for task, names in self._names.items(): + sums = self._sums[task] + weights = self._weights[task] + for name in names: + sums[name] = values[index] + index += 1 + for name in names: + weights[name] = values[index] + index += 1 + self._observations[task] = int(float(values[index])) + index += 1 + + +# Historical name kept for call sites and tests that predate the weighted API. +TrainingMetricAccumulator = MetricAccumulator diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 531bc7d07c..6901077b8b 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -56,7 +56,7 @@ if TYPE_CHECKING: from .metrics import ( - TrainingMetricAccumulator, + MetricAccumulator, ) DEFAULT_TASK_KEY = "Default" @@ -537,7 +537,7 @@ def __init__( trainer_config: TrainerConfig, *, rank_context: RankContext | None = None, - metric_accumulator: TrainingMetricAccumulator | None = None, + metric_accumulator: MetricAccumulator | None = None, ) -> None: self.trainer_config = trainer_config self.rank_context = rank_context or RankContext() @@ -572,6 +572,10 @@ def run(self, tasks: TrainingTaskCollection) -> None: display_step = step + 1 if self._should_display(display_step): + # Collectives for averaged metrics run on every rank before + # the chief-only logging path. Local windows stay on device + # until this boundary. + self.synchronize_metric_accumulator() if self.rank_context.is_chief: train_results, valid_results = self.collect_display_results( tasks, @@ -665,6 +669,16 @@ def collect_display_results( return train_results[active_task.key], valid_results[active_task.key] return train_results, valid_results + def synchronize_metric_accumulator(self) -> None: + """Reduce local metric windows across ranks at a display boundary. + + The default implementation is a no-op for single-process runs. A + distributed backend overrides this with an all-reduce of the packed + sums, weights, and observation counts. The call is collective: every + rank that accumulates metrics must enter it together. + """ + return None + def on_train_begin(self, tasks: TrainingTaskCollection) -> None: """Hook called before the first optimizer step.""" return None diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7a7210962e..285a113e0d 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -28,6 +28,7 @@ from deepmd.dpmodel.train import ( DEFAULT_TASK_KEY, CheckpointStore, + MetricAccumulator, ShardingPolicy, TrainingMetricAccumulator, TrainingTimer, @@ -119,6 +120,9 @@ NonFiniteGradGuard, clip_grad_norm_, ) +from deepmd.pt_expt.train.metrics import ( + all_reduce_metric_accumulator, +) from deepmd.pt_expt.train.utils import ( count_parameters, infer_env_defaults, @@ -1593,22 +1597,28 @@ def fake_model() -> dict: display_step_id % self.disp_freq == 0 or display_step_id == 1 ): self.wrapper.eval() # Will set to train mode before fininshing validation + if self.metric_accumulator is not None and self.world_size > 1: + all_reduce_metric_accumulator(self.metric_accumulator) def log_loss_valid(_task_key: str = "Default") -> dict: - single_results = {} - sum_natoms = 0 if not self.multi_task: valid_numb_batch = self.valid_numb_batch + loss_obj = self.loss else: valid_numb_batch = self.valid_numb_batch[_task_key] + loss_obj = self.loss[_task_key] + accumulator = MetricAccumulator( + {_task_key: loss_obj.training_metric_names} + ) + saw_batch = False for ii in range(valid_numb_batch): self.optimizer.zero_grad() input_dict, label_dict, _ = self.get_data( is_train=False, task_key=_task_key ) if input_dict == {}: - # no validation data - return {} + # no validation data for this slot + break # Validation runs the inner module, not the DDP # wrapper. A DDP forward under grad mode arms the # reducer for an all-reduce that this loop never @@ -1624,20 +1634,24 @@ def log_loss_valid(_task_key: str = "Default") -> dict: label=label_dict, task_key=_task_key, ) - # more_loss.update({"rmse": math.sqrt(loss)}) - # The metrics are per-atom quantities, so each batch - # weighs by the real atoms it holds summed over its - # frames. Phantom atoms (atype < 0), which pad a - # mixed-nloc batch, contribute to none of them. + # Per-atom metrics weigh by real atoms across frames. + # Phantom atoms (atype < 0) pad mixed-nloc batches. natoms = int((input_dict["atype"] >= 0).sum()) - sum_natoms += natoms - for k, v in more_loss.items(): - if "l2_" not in k: - single_results[k] = ( - single_results.get(k, 0.0) + v * natoms + accumulator.add( + _task_key, + { + name: ( + value.detach() if torch.is_tensor(value) else value ) - results = {k: v / sum_natoms for k, v in single_results.items()} - return results + for name, value in more_loss.items() + if "l2_" not in name + }, + weight=float(natoms), + ) + saw_batch = True + if not saw_batch: + return {} + return accumulator.average(_task_key) if not self.multi_task: train_results = self._training_results(more_loss) diff --git a/deepmd/pt_expt/train/metrics.py b/deepmd/pt_expt/train/metrics.py new file mode 100644 index 0000000000..d204153a1a --- /dev/null +++ b/deepmd/pt_expt/train/metrics.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Torch helpers for distributed metric-window reduction.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from deepmd.dpmodel.train.metrics import ( + MetricAccumulator, + ) + + +def all_reduce_metric_accumulator( + accumulator: MetricAccumulator, + *, + group: Any | None = None, +) -> None: + """Sum local metric windows across ranks in place. + + Call only at a display boundary. Host synchronization of the averaged + scalars still happens later in :meth:`MetricAccumulator.average`; this + routine only exchanges the packed sums, weights, and observation counts. + Single-process runs leave the accumulator unchanged. + """ + if not dist.is_available() or not dist.is_initialized(): + return + world_size = dist.get_world_size(group) + if world_size <= 1: + return + + packed = accumulator.flat_state() + if not packed: + return + + # Packing always host-converts via float(); the collective device must + # come from the process group, not from whether this rank's window + # promoted any sums to tensors. + device = _collective_device(group) + tensor = torch.tensor( + [float(value) for value in packed], + dtype=torch.float64, + device=device, + ) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group) + accumulator.load_flat_state(tensor.detach().cpu().tolist()) + + +def _collective_device(group: Any | None = None) -> torch.device: + """Return one all-reduce device that every rank in ``group`` can use. + + Device choice follows the process-group backend and this rank's current + CUDA device. It must not depend on local accumulator contents: a rank + whose window never promoted sums to tensors would otherwise pick CPU + while another rank picked CUDA, which breaks NCCL collectives. + """ + backend = str(dist.get_backend(group)).lower() + if "nccl" in backend and torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..fb8da802b5 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -49,6 +49,7 @@ from deepmd.dpmodel.train import ( DEFAULT_TASK_KEY, AbstractTrainer, + MetricAccumulator, RankContext, ShardingPolicy, TrainerConfig, @@ -120,6 +121,9 @@ NonFiniteGradGuard, clip_grad_norm_, ) +from deepmd.pt_expt.train.metrics import ( + all_reduce_metric_accumulator, +) from deepmd.pt_expt.train.utils import ( MatmulPrecisionPolicy, count_parameters, @@ -3278,6 +3282,12 @@ def on_train_begin(self, tasks: TrainingTaskCollection) -> None: """Switch the wrapper to training mode.""" self.wrapper.train() + def synchronize_metric_accumulator(self) -> None: + """All-reduce averaged training metrics before the chief logs them.""" + if self.metric_accumulator is None or self.world_size <= 1: + return + all_reduce_metric_accumulator(self.metric_accumulator) + def collect_display_results( self, tasks: TrainingTaskCollection, @@ -3382,12 +3392,17 @@ def evaluate_validation( Validation is therefore skipped from stage two on; the metrics remain available through the independent full validation flow, which every rank enters together. + + Display validation uses the shared :class:`MetricAccumulator` with + atom weights so optional NaN metrics and empty batches share the same + rules as training interval averages. """ if task.validation_data is None or self.sharding.shards_parameters: return None - valid_results: dict[str, float] = {} - sum_natoms = 0 + loss = self.losses[task.key] + accumulator = MetricAccumulator({task.key: loss.training_metric_names}) + saw_batch = False for _ii in range(task.valid_numb_batch): val_input, val_label = self.get_data(is_train=False, task_key=task.key) if not val_input: @@ -3398,21 +3413,22 @@ def evaluate_validation( label=val_label, task_key=task.key, ) - # The metrics are per-atom quantities, so each batch weighs by the - # real atoms it holds summed over its frames. Phantom atoms - # (atype < 0), which pad a mixed-nloc batch, contribute to none. + # Per-atom metrics weigh by real atoms across frames. Phantom + # atoms (atype < 0) pad mixed-nloc batches and contribute nothing. natoms = int((val_input["atype"] >= 0).sum()) - sum_natoms += natoms - for key, value in vmore.items(): - if "l2_" not in key: - valid_results[key] = ( - valid_results.get(key, 0.0) + self._to_float(value) * natoms - ) - if sum_natoms > 0: - valid_results = { - key: value / sum_natoms for key, value in valid_results.items() - } - return valid_results + accumulator.add( + task.key, + { + key: value.detach() if torch.is_tensor(value) else value + for key, value in vmore.items() + if "l2_" not in key + }, + weight=float(natoms), + ) + saw_batch = True + if not saw_batch: + return {} + return accumulator.average(task.key) def learning_rate(self, step: int) -> float: """Return the configured learning rate for a zero-based step.""" diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 982c5d2a26..f4e748c167 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6023,8 +6023,12 @@ def training_args( doc_disp_avg = ( "Display the arithmetic mean of per-step training metrics within each " "display interval, separately for each task. Tasks with no training " - "steps in an interval display NaN. Validation metrics are evaluated " - "at the display step, not averaged over training steps." + "steps in an interval display NaN. A metric that is NaN on a given " + "step (for example an optional label that was absent) is excluded " + "from that metric's interval mean. Validation metrics are evaluated " + "at the display step with atom weighting, not averaged over training " + "steps. Distributed runs reduce local sums and weights only at " + "display boundaries." ) doc_profiling = ( "Enable performance profiling. TensorFlow and PyTorch can export a Chrome " diff --git a/source/tests/common/dpmodel/test_train_metrics.py b/source/tests/common/dpmodel/test_train_metrics.py index 8acabcae7a..ab6a34a450 100644 --- a/source/tests/common/dpmodel/test_train_metrics.py +++ b/source/tests/common/dpmodel/test_train_metrics.py @@ -7,6 +7,7 @@ import pytest from deepmd.dpmodel.train.metrics import ( + MetricAccumulator, TrainingMetricAccumulator, ) @@ -45,14 +46,91 @@ def test_accumulation_does_not_alias_or_modify_backend_buffers() -> None: @pytest.mark.parametrize("metrics", [{}, {"rmse": float("nan")}]) -def test_missing_and_nan_metrics_are_not_silently_zeroed(metrics: dict) -> None: - accumulator = TrainingMetricAccumulator({"task": ("rmse",)}) +def test_missing_and_nan_metrics_are_excluded_from_the_average(metrics: dict) -> None: + """Optional labels must not poison an otherwise finite interval average.""" + accumulator = MetricAccumulator({"task": ("rmse",)}) accumulator.add("task", {"rmse": 1.0}) accumulator.add("task", metrics) + assert accumulator.average("task") == {"rmse": 1.0} + assert accumulator.count("task") == 2 + + +def test_metric_with_only_nan_observations_reports_nan() -> None: + accumulator = MetricAccumulator({"task": ("rmse", "mae")}) + accumulator.add("task", {"rmse": float("nan"), "mae": 2.0}) + accumulator.add("task", {"rmse": float("nan"), "mae": 4.0}) assert math.isnan(accumulator.average("task")["rmse"]) + assert accumulator.average("task")["mae"] == 3.0 + + +def test_add_does_not_float_coerce_device_scalars() -> None: + """Present accelerator metrics must not D2H-sync on every training step.""" + + class DeviceScalar: + def __init__(self, value: float) -> None: + self.value = float(value) + self.device = "cuda:0" + self.ndim = 0 + self.float_calls = 0 + + def __float__(self) -> float: + self.float_calls += 1 + return self.value + + def __mul__(self, other: object): + return DeviceScalar(self.value * float(other)) + + __rmul__ = __mul__ + + def __add__(self, other: object): + if isinstance(other, DeviceScalar): + return DeviceScalar(self.value + other.value) + return DeviceScalar(self.value + float(other)) + + def __radd__(self, other: object): + return self.__add__(other) + + def __truediv__(self, other: object): + return DeviceScalar(self.value / float(other)) + + accumulator = MetricAccumulator({"task": ("rmse",)}) + device_value = DeviceScalar(2.0) + accumulator.add("task", {"rmse": device_value}) + assert device_value.float_calls == 0 + # Host conversion is reserved for the display boundary. + assert accumulator.average("task") == {"rmse": 2.0} + + +def test_atom_weights_compute_validation_style_averages() -> None: + accumulator = MetricAccumulator({"task": ("rmse",)}) + accumulator.add("task", {"rmse": 1.0}, weight=2.0) + accumulator.add("task", {"rmse": 4.0}, weight=1.0) + assert accumulator.average("task") == {"rmse": 2.0} assert accumulator.count("task") == 2 +def test_non_positive_weight_skips_metric_totals() -> None: + accumulator = MetricAccumulator({"task": ("rmse",)}) + accumulator.add("task", {"rmse": 1.0}, weight=0.0) + assert accumulator.count("task") == 1 + assert math.isnan(accumulator.average("task")["rmse"]) + + +def test_flat_state_round_trip_supports_distributed_reduce() -> None: + left = MetricAccumulator({"a": ("rmse",), "b": ("mae",)}) + right = MetricAccumulator({"a": ("rmse",), "b": ("mae",)}) + left.add("a", {"rmse": 1.0}) + left.add("b", {"mae": 2.0}) + right.add("a", {"rmse": 3.0}) + # b unsampled on the right rank + merged = [x + y for x, y in zip(left.flat_state(), right.flat_state(), strict=True)] + left.load_flat_state(merged) + assert left.average("a") == {"rmse": 2.0} + assert left.average("b") == {"mae": 2.0} + assert left.count("a") == 2 + assert left.count("b") == 1 + + def test_unknown_metric_does_not_partially_update_the_window() -> None: accumulator = TrainingMetricAccumulator({"task": ("rmse",)}) with pytest.raises(ValueError, match=r"Undeclared display metrics.*mae"): diff --git a/source/tests/pt_expt/train/test_metric_reduce.py b/source/tests/pt_expt/train/test_metric_reduce.py new file mode 100644 index 0000000000..5c1d07921c --- /dev/null +++ b/source/tests/pt_expt/train/test_metric_reduce.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for Torch metric-window reduction helpers.""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import pytest + +torch = pytest.importorskip("torch") + +from deepmd.dpmodel.train.metrics import ( + MetricAccumulator, +) +from deepmd.pt_expt.train import metrics as metrics_mod +from deepmd.pt_expt.train.metrics import ( + _collective_device, + all_reduce_metric_accumulator, +) + + +def test_all_reduce_is_noop_without_process_group() -> None: + accumulator = MetricAccumulator({"task": ("rmse",)}) + accumulator.add("task", {"rmse": 1.5}) + all_reduce_metric_accumulator(accumulator) + assert accumulator.average("task") == {"rmse": 1.5} + + +def test_collective_device_follows_nccl_not_host_window( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """NCCL ranks must agree on CUDA even when the local window is host-only.""" + monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "nccl") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + assert _collective_device() == torch.device("cuda", 0) + + monkeypatch.setattr( + metrics_mod.dist, + "get_backend", + lambda group=None: "cuda:nccl,cpu:gloo", + ) + assert _collective_device() == torch.device("cuda", 0) + + +def test_collective_device_uses_cpu_for_gloo(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "gloo") + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + assert _collective_device() == torch.device("cpu") + + +def test_all_reduce_uses_group_device_for_host_only_window( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Host-only flat_state must still all-reduce on the group device.""" + captured: dict[str, Any] = {} + + def fake_all_reduce(tensor: torch.Tensor, op=None, group=None) -> None: + captured["device"] = tensor.device + tensor.fill_(0.0) + + monkeypatch.setattr(metrics_mod.dist, "is_available", lambda: True) + monkeypatch.setattr(metrics_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(metrics_mod.dist, "get_world_size", lambda group=None: 2) + monkeypatch.setattr(metrics_mod.dist, "get_backend", lambda group=None: "nccl") + monkeypatch.setattr(metrics_mod.dist, "all_reduce", fake_all_reduce) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + # Keep tensor construction on CPU in this CPU-only CI image while still + # asserting that the requested device is the NCCL CUDA device. + real_tensor = torch.tensor + + def fake_tensor(data, dtype=None, device=None, **kwargs): + captured["requested_device"] = ( + torch.device(device) if device is not None else torch.device("cpu") + ) + return real_tensor(data, dtype=dtype, device="cpu", **kwargs) + + monkeypatch.setattr(torch, "tensor", fake_tensor) + + accumulator = MetricAccumulator({"task": ("rmse",)}) + # Optional-label NaNs leave sums as host floats (never promoted to CUDA). + accumulator.add("task", {"rmse": float("nan")}) + assert all(not torch.is_tensor(v) for v in accumulator.flat_state()) + + all_reduce_metric_accumulator(accumulator) + assert captured["requested_device"] == torch.device("cuda", 0)