Update Wall-X to 1.1.0 (#104)

wall_x/trainer/fsdp_trainer/metrics/norms.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Callable, DefaultDict, Optional
+
+import torch
+from torch.distributed.tensor import DTensor, Replicate
+
+
+class DistributedNormAccumulator:
+    """Accumulate L2 norms without leaking parallel-layout branches to adapters."""
+
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        reduce_tensor_fn: Optional[Callable[[torch.Tensor, bool], torch.Tensor]],
+    ):
+        self.device = device
+        self.reduce_tensor_fn = reduce_tensor_fn
+        self._regular: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
+        self._dtensor: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
+        self._owner_only: DefaultDict[str, list[torch.Tensor]] = defaultdict(list)
+
+    @staticmethod
+    def _is_dtensor(tensor: torch.Tensor) -> bool:
+        return isinstance(tensor, DTensor)
+
+    def add(self, group: Optional[str], tensor: Optional[torch.Tensor]) -> None:
+        if group is None or tensor is None or tensor.numel() == 0:
+            return
+        tensor = tensor.detach()
+        if self._is_dtensor(tensor):
+            self._dtensor[group].append(tensor)
+        else:
+            self._regular[group].append(tensor)
+
+    def add_owner_only(
+        self, group: Optional[str], tensor: Optional[torch.Tensor]
+    ) -> None:
+        if group is None or tensor is None or tensor.numel() == 0:
+            return
+        self._owner_only[group].append(tensor.detach())
+
+    def squared_norm(self, group: str) -> torch.Tensor:
+        with torch.no_grad():
+            total = torch.zeros((), dtype=torch.float32, device=self.device)
+            total = total + self._regular_squared_norm(self._regular[group])
+            total = total + self._owner_only_squared_norm(self._owner_only[group])
+            total = total + self._dtensor_squared_norm(self._dtensor[group])
+            return total
+
+    def norm(self, group: str) -> torch.Tensor:
+        return self.squared_norm(group).sqrt()
+
+    def _regular_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
+        if not tensors:
+            return torch.zeros((), dtype=torch.float32, device=self.device)
+        per = self._foreach_norm(tensors)
+        return self._sum_norms_squared(per)
+
+    def _owner_only_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
+        local_sq = self._regular_squared_norm(tensors)
+        if self.reduce_tensor_fn is None:
+            return local_sq
+        return self.reduce_tensor_fn(local_sq, average=False)
+
+    def _dtensor_squared_norm(self, tensors: list[torch.Tensor]) -> torch.Tensor:
+        if not tensors:
+            return torch.zeros((), dtype=torch.float32, device=self.device)
+
+        per = self._foreach_norm(tensors)
+
+        # Group by (mesh, placements) so each unique layout costs only one
+        # collective (redistribute).
+        groups: DefaultDict[object, list[torch.Tensor]] = defaultdict(list)
+        for norm in per:
+            if self._is_dtensor(norm):
+                key = (id(norm.device_mesh), tuple(norm.placements))
+            else:
+                key = None
+            groups[key].append(norm)
+
+        parts: list[torch.Tensor] = []
+        for group_tensors in groups.values():
+            try:
+                stacked = torch.stack(group_tensors)
+                local = self._to_local_replicated(stacked).to(
+                    device=self.device, dtype=torch.float32
+                )
+                parts.append(local.pow(2).sum())
+            except (RuntimeError, NotImplementedError):
+                # Stack/redistribute not supported for this layout in the
+                # current PyTorch version; fall back to per-tensor for safety.
+                for norm in group_tensors:
+                    local = self._to_local_replicated(norm).to(
+                        device=self.device, dtype=torch.float32
+                    )
+                    parts.append(local.pow(2).sum())
+
+        return torch.stack(parts).sum()
+
+    def _foreach_norm(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]:
+        try:
+            return torch._foreach_norm(tensors, 2.0)
+        except (RuntimeError, NotImplementedError):
+            return [torch.linalg.vector_norm(t, 2.0) for t in tensors]
+
+    def _sum_norms_squared(self, norms: list[torch.Tensor]) -> torch.Tensor:
+        if not norms:
+            return torch.zeros((), dtype=torch.float32, device=self.device)
+        local_norms = [
+            self._to_local_replicated(norm).to(device=self.device, dtype=torch.float32)
+            for norm in norms
+        ]
+        return torch.stack(local_norms).pow(2).sum()
+
+    def _to_local_replicated(self, tensor: torch.Tensor) -> torch.Tensor:
+        if not self._is_dtensor(tensor):
+            return tensor.to(device=self.device)
+        mesh = tensor.device_mesh
+        tensor = tensor.redistribute(placements=[Replicate()] * mesh.ndim)
+        return tensor.to_local().to(device=self.device)