Sync 1 diff to facebookresearch/optimizers

CHANGELOG.md

@@ -1,6 +1,16 @@
 # Changelog
 
-Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the internal source tree.
+Release notes for the OSS mirror. Seeded by `scripts/oss_sync.sh` from the internal source tree and edited for release.
+
+## 2026-08-21
+
+- [shampoo] Add per-head block splitting for attention projections, via two new `DDPDistributedConfig` fields: `head_split_size` and `head_split_param_names`. When `head_split_size > 0`, a matching 2D parameter is split along dim 0 into `head_split_size`-row blocks instead of taking the default merge-and-block path, so a matrix preconditioner (Shampoo or spectral descent) sees one block per attention head rather than the whole fused projection. A parameter is split iff its name contains any substring in `head_split_param_names`; parameters in the same group that do not match keep the default blocking. Opt-in — `head_split_size` defaults to 0, so the default path is unchanged. Requires per-parameter names under `param_names` in the param group, and raises if they are missing or their count does not match `params`; a negative `head_split_size` raises at config construction, and at block time a non-2D parameter, or a size that does not divide the parameter's first dimension, raises. If `head_split_size` is set but no name matches, a warning is logged and the split is a no-op. Implemented for the DDP distributor.
+
+## 2026-08-19
+
+- [shampoo] Add `NewtonSchulzRootInvConfig`: a matmul-only coupled Newton-Schulz iteration, usable as the `amortized_computation_config` of `RootInvShampooPreconditionerConfig` in place of the default eigendecomposition. Opt-in — `DefaultShampooConfig` is unchanged. The iteration runs a fixed number of steps with no residual-based stopping criterion, so it introduces no host-device synchronization. Only power-of-two inverse roots are supported (blocks of order 1, 2, or 4); anything else raises at optimizer construction rather than degrading to a per-factor-matrix warning that reuses a stale preconditioner. `relative_epsilon` (default `1e-6`) floors the ridge at a fraction of `|A|_F` and is applied unconditionally, so on the rank-deficient factor matrices seen early in training this path regularizes more aggressively than eigendecomposition at the same epsilon and the two do not agree there. `coefficients` takes a per-iteration schedule of `(a, b, c)` triples for `p(x) = a x + b x^3 + c x^5`, defaulting to a 10-step Polar Express schedule; tf32 is disabled inside the iteration by default (`disable_tf32=True`).
+- [shampoo] Annotate `ShampooPT2CompileConfig`'s `make_dataclass`-synthesized base class `Any` so mypy accepts it as a base class, and add test coverage asserting the synthesized field set matches `torch.compile`'s signature (no runtime change).
+- [shampoo] Remove unused type-error suppression comments across the optimizer, examples, and tests (no behavior change).
 
 ## 2026-07-24
 

distributed_shampoo/__init__.py

@@ -23,6 +23,7 @@
     EighEigendecompositionConfig,
     MatrixFunctionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -118,6 +119,7 @@
     "DefaultEigenConfig",  # Default `RootInvConfig` using `EigenConfig`.
     "CoupledNewtonConfig",  # Based on `RootInvConfig`.
     "CoupledHigherOrderConfig",  # Based on `RootInvConfig`.
+    "NewtonSchulzRootInvConfig",  # Based on `RootInvConfig`.
     "OrthogonalizationConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
     "SVDOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.
     "NewtonSchulzOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.

distributed_shampoo/distributed_shampoo.py

@@ -981,7 +981,6 @@ def _instantiate_per_group_step(
         # Use PT2 to compile the step function for each parameter group.
         self._per_group_step: Callable[..., None] = (
             torch.compile(
-                # pyrefly: ignore [bad-argument-type]
                 self._per_group_step_impl,
                 # pyrefly: ignore [bad-argument-type]
                 **asdict(shampoo_pt2_compile_config),
@@ -1370,7 +1369,7 @@ def _apply_in_place_primal_averaging(
         # This computes: - (1 - mu_x * mu_y) * lr * P.
         torch._foreach_mul_(
             masked_blocked_search_directions,
-            (1 - train_interp_coeff * eval_interp_coeff),  # type: ignore
+            (1 - train_interp_coeff * eval_interp_coeff),
         )
         # This computes: (1 - mu_x) * (Z_old - Y) - (1 - mu_x * mu_y) * lr * P.
         # pyrefly: ignore [no-matching-overload]

distributed_shampoo/distributor/shampoo_ddp_distributor.py

@@ -24,6 +24,7 @@
 from distributed_shampoo.shampoo_types import (
     DDPDistributedConfig,
     DISTRIBUTED_CONFIG,
+    PARAM_NAMES,
     PARAMS,
     ShampooRuntimeConfig,
 )
@@ -220,6 +221,47 @@ def __init__(
             group_rank=group_rank,
         )
 
+    def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None:
+        """Resolve per-parameter head-split sizes from the DDP config.
+
+        Returns None (head-split disabled) unless ``head_split_size > 0``. When
+        enabled, a parameter is head-split iff its name contains any substring in
+        ``head_split_param_names``; ``head_split_size`` (head_dim) is used for
+        those and None for the rest. Requires per-parameter names under
+        PARAM_NAMES in the param group. Works with any matrix preconditioner
+        (e.g. Shampoo or spectral descent / Muon).
+        """
+        distributed_config: DDPDistributedConfig = self._param_group[DISTRIBUTED_CONFIG]
+        head_split_size = distributed_config.head_split_size
+        if not head_split_size:
+            return None
+        if PARAM_NAMES not in self._param_group:
+            raise ValueError(
+                f"head_split_size requires per-parameter names under "
+                f"'{PARAM_NAMES}' in the param group."
+            )
+        param_names = self._param_group[PARAM_NAMES]
+        num_params = len(self._param_group[PARAMS])
+        if len(param_names) != num_params:
+            raise ValueError(
+                f"len({PARAM_NAMES})={len(param_names)} must equal "
+                f"len(params)={num_params}."
+            )
+        head_split_sizes = tuple(
+            head_split_size
+            if any(
+                pattern in name for pattern in distributed_config.head_split_param_names
+            )
+            else None
+            for name in param_names
+        )
+        if not any(size is not None for size in head_split_sizes):
+            logger.warning(
+                "head_split_size is set but no parameter name matched "
+                "head_split_param_names; head-split is a no-op."
+            )
+        return head_split_sizes
+
     @overload
     @torch.no_grad()
     def _get_params_or_grads(

distributed_shampoo/distributor/shampoo_distributor.py

@@ -7,8 +7,9 @@
 
 """
 
+import logging
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from functools import partial
 from operator import attrgetter
 from typing import Any, Literal, overload
@@ -29,6 +30,8 @@
 )
 from torch import Tensor
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 ###### DISTRIBUTOR CLASSES ######
 class DistributorInterface(ABC):
@@ -51,6 +54,13 @@ def __init__(
         self._runtime_config: ShampooRuntimeConfig = (
             runtime_config if runtime_config is not None else ShampooRuntimeConfig()
         )
+        # Per-parameter head-split sizes (None disables head-split for every
+        # param). Resolved via virtual dispatch so only distributors that
+        # support head-split (currently DDP) return a non-None spec. Must be set
+        # before _merge_and_block_parameters() so blocking can consult it.
+        self._head_split_sizes: tuple[int | None, ...] | None = (
+            self._resolve_head_split_sizes()
+        )
         # Merge and block parameters creates self._global_blocked_params and self._global_num_blocks_per_param
         # Global blocked params are all the blocked parameters after merging and blocking.
         # Global num blocks per param stores the number of blocks for each global parameter.
@@ -75,6 +85,21 @@ def __init__(
         # Local block info list contains information about each block masked by the distributor selector.
         self._local_block_info_list: tuple[BlockInfo, ...]
 
+    def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None:
+        """Return per-parameter head-split sizes, or None to disable head-split.
+
+        Base implementation disables head-split. Distributors that support it
+        (currently only ``DDPDistributor``) override this. A warning is emitted
+        if head-split was requested on an unsupported distributor so the config
+        is not silently ignored.
+        """
+        if getattr(self._param_group[DISTRIBUTED_CONFIG], "head_split_size", 0):
+            logger.warning(
+                "head_split_size is set but ignored: head-split is only "
+                f"supported by DDPDistributor, not {type(self).__name__}."
+            )
+        return None
+
     @abstractmethod
     @torch.no_grad()
     def update_params(
@@ -170,6 +195,43 @@ def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None
             else self._param_group[PARAMS]
         )
 
+    def _blocks_within_tensor(
+        self,
+        tensor: Tensor,
+        param_index: int,
+        merge_dims: Callable[..., tuple[int, ...]],
+    ) -> tuple[Tensor, ...]:
+        """Split one param/grad tensor into preconditioner blocks.
+
+        Default path: merge small dims, then block by ``max_preconditioner_dim``.
+        When head-split is enabled for this param, instead split dim 0 into
+        ``head_dim``-sized 2D blocks (one per attention head) and skip merging /
+        blocking so the model (input) dimension is preserved. Both the parameter
+        and gradient paths call this so their block decomposition stays
+        identical.
+        """
+        head_dim = (
+            None
+            if self._head_split_sizes is None
+            else self._head_split_sizes[param_index]
+        )
+        if head_dim is not None:
+            if tensor.dim() != 2:
+                raise ValueError(
+                    "head-split requires a 2D parameter, got shape "
+                    f"{tuple(tensor.shape)} at param index {param_index}."
+                )
+            if tensor.shape[0] % head_dim != 0:
+                raise ValueError(
+                    f"head_split_size={head_dim} must divide out_features="
+                    f"{tensor.shape[0]} at param index {param_index}."
+                )
+            return torch.split(tensor, head_dim, dim=0)
+        return multi_dim_split(
+            tensor.view(merge_dims(tensor_shape=tensor.size())),
+            self._param_group[MAX_PRECONDITIONER_DIM],
+        )
+
     @torch.no_grad()
     def _merge_and_block_with_params(
         self, params: Iterable[Tensor]
@@ -199,11 +261,10 @@ def _merge_and_block_with_params(
             ].target_parameter_dimensionality,
         )
 
-        for param in params:
-            # Obtain blocks for each parameter after merging.
-            blocks_within_param = multi_dim_split(
-                param.view(merge_dims(tensor_shape=param.size())),
-                self._param_group[MAX_PRECONDITIONER_DIM],
+        for param_index, param in enumerate(params):
+            # Obtain blocks for each parameter after merging (or per-head split).
+            blocks_within_param = self._blocks_within_tensor(
+                param, param_index, merge_dims
             )
 
             # Generate and extend blocked parameters list.
@@ -256,11 +317,17 @@ def _merge_and_block_gradients(
             ].target_parameter_dimensionality,
         )
 
-        for grad, num_blocks, (block_index, next_block_index) in zip(
-            self._get_params_or_grads(get_grad=True),
-            self._global_num_blocks_per_param,
-            generate_pairwise_indices(self._global_num_blocks_per_param),
-            strict=True,
+        for param_index, (
+            grad,
+            num_blocks,
+            (block_index, next_block_index),
+        ) in enumerate(
+            zip(
+                self._get_params_or_grads(get_grad=True),
+                self._global_num_blocks_per_param,
+                generate_pairwise_indices(self._global_num_blocks_per_param),
+                strict=True,
+            )
         ):
             param_distributor_selector = self._distributor_selector[
                 block_index:next_block_index
@@ -283,10 +350,9 @@ def _merge_and_block_gradients(
                     f"Encountered gradient containing NaN/Inf in parameter with shape {attrgetter('shape')(grad)}. Check your model for numerical instability or consider gradient clipping."
                 )
 
-            # Obtain blocks for each gradient after merging.
-            blocks_within_grad = multi_dim_split(
-                grad.view(merge_dims(tensor_shape=grad.size())),
-                self._param_group[MAX_PRECONDITIONER_DIM],
+            # Obtain blocks for each gradient after merging (or per-head split).
+            blocks_within_grad = self._blocks_within_tensor(
+                grad, param_index, merge_dims
             )
             # Generate block-to-parameter metadata and extend blocked parameters list.
             local_masked_blocked_grads.extend(

distributed_shampoo/examples/utils.py

@@ -73,7 +73,6 @@ def get_distributed_env() -> tuple[int, int, int]:
 
 def set_seed(seed: int) -> None:
     torch.manual_seed(seed)
-    # pyrefly: ignore [bad-argument-type]
     np.random.seed(seed)
     random.seed(seed)
     torch.use_deterministic_algorithms(True)