Lightning-AI/pytorch-lightning · error · ValueError

Currently, the XLAStrategy only supports `sum`, `mean`, `avg

Error message

Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got: {reduce_op}

What it means

XLAStrategy.reduce only supports SUM-like reductions because it maps onto torch_xla's xm.reduce_type / all-reduce, which supports 'sum', 'mean', 'avg' (and ReduceOp.SUM). Requesting other ops (MAX, MIN, PRODUCT, etc.) raises ValueError.

Source

Thrown at src/lightning/pytorch/strategies/xla.py:261

        else:
            obj = obj.to(original_device)

        return obj

    @override
    def reduce(
        self,
        output: Union[Tensor, Any],
        group: Optional[Any] = None,
        reduce_op: Optional[Union[ReduceOp, str]] = "mean",
    ) -> Tensor:
        if not isinstance(output, Tensor):
            output = torch.tensor(output, device=self.root_device)

        invalid_reduce_op = isinstance(reduce_op, ReduceOp) and reduce_op != ReduceOp.SUM
        invalid_reduce_op_str = isinstance(reduce_op, str) and reduce_op.lower() not in ("sum", "mean", "avg")
        if invalid_reduce_op or invalid_reduce_op_str:
            raise ValueError(
                "Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got:"
                f" {reduce_op}"
            )

        import torch_xla.core.xla_model as xm

        output = xm.mesh_reduce("reduce", output, sum)

        if isinstance(reduce_op, str) and reduce_op.lower() in ("avg", "mean"):
            output = output / self.world_size

        return output

    @override
    def setup_environment(self) -> None:
        self._launched = True
        super().setup_environment()

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Compute sum-based reductions instead (e.g. transform max into -min of negatives: reduce(-x, 'min-free' sum trick)) or implement max via all_gather + local max
2. Use 'sum'/'mean'/'avg' only when calling strategy.reduce
3. Do the min/max locally after an all_gather of values

Example fix

# before
val = self.trainer.strategy.reduce(tensor, reduce_op="max")  # ValueError

# after
gathered = self.trainer.strategy.all_gather(tensor)
val = gathered.max(dim=0).values  # compute max after gather

op = str(reduce_op).lower()
assert op in ("sum", "mean", "avg") or reduce_op in (None, ReduceOp.SUM), "XLA reduce supports only sum/mean/avg"

def xla_reduce_ok(reduce_op) -> bool:
    import torch.distributed as dist
    if isinstance(reduce_op, dist.ReduceOp):
        return reduce_op == dist.ReduceOp.SUM
    return isinstance(reduce_op, str) and reduce_op.lower() in ("sum", "mean", "avg")

When it happens

Trigger: Calling strategy.reduce(tensor, reduce_op='max') or ReduceOp.MAX / any op other than sum/mean/avg; using LightningModule.all_gather/merge things or custom code invoking reduce with unsupported ops under XLAStrategy.

Common situations: Custom metrics that reduce max/min across ranks; porting DDP code where torch.distributed.ReduceOp.MAX works fine.

Related errors

• `num_nodes` must be a positive integer, but got {num_nodes}.
• To spawn processes with the `{type(self.strategy).__name__}`
• To use Fabric with more than one device, you must call `.lau
• The `{type(self._strategy).__name__}` requires the model and
• `{type(self).__name__}` does not own a group. HINT: try `col

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/1e0d72f66ed4d01c. Report an issue: GitHub.