diff --git a/distributed_shampoo/README.md b/distributed_shampoo/README.md
index c22356b..13b9434 100644
--- a/distributed_shampoo/README.md
+++ b/distributed_shampoo/README.md
@@ -21,7 +21,7 @@ Developers:
 
 with contributions and support from:
 
-Ganesh Ajjanagadde (Meta), Rohan Anil (Google), Adnan Aziz (Meta), Pavan Balaji (Meta), Shuo Chang (Meta), Weiwei Chu (Meta), Assaf Eisenman (Meta), Will Feng (Meta), Zhuobo Feng (Meta), Jose Gallego-Posada (Mila / Meta Platforms, Inc.), Avirup Ghosh (Meta), Yizi Gu (Meta), Vineet Gupta (Google), Yuchen Hao (Meta), Brian Hirsh (Meta), Yusuo Hu (Meta), Yuxi Hu (Meta), Minhui Huang (Meta), Guna Lakshminarayanan (Meta), Michael Lazos (Meta), Zhijing Li (Meta), Ming Liang (Meta), Wanchao Liang (Meta), Ying Liu (Meta), Wenguang Mao (Meta), Dheevatsa Mudigere (NVIDIA), Maxim Naumov (Meta), Jongsoo Park (Meta), Mike Rabbat (Meta), Kaushik Rangadurai (Meta), Dennis van der Staay (Meta), Fei Tian (Meta), Rohan Varma (Meta), Sanjay Vishwakarma (Meta), Xunnan (Shawn) Xu (Meta), Jiyan Yang (Meta), Chunxing Yin (Meta), Gavin Zhang (Meta), Haoran Zhang (Meta), Haoyu Zhang (Meta), Chuanhao Zhuge (Meta), and Will Zou (Meta).
+Ganesh Ajjanagadde (Meta), Rohan Anil (Google), Adnan Aziz (Meta), Pavan Balaji (Meta), Shuo Chang (Meta), Weiwei Chu (Meta), Assaf Eisenman (Meta), Will Feng (Meta), Zhuobo Feng (Meta), Jose Gallego-Posada (Mila / Meta Platforms, Inc.), Avirup Ghosh (Meta), Yizi Gu (Meta), Shagun Gupta (Meta), Vineet Gupta (Google), Yuchen Hao (Meta), Brian Hirsh (Meta), Yusuo Hu (Meta), Yuxi Hu (Meta), Minhui Huang (Meta), Guna Lakshminarayanan (Meta), Michael Lazos (Meta), Zhijing Li (Meta), Ming Liang (Meta), Wanchao Liang (Meta), Ying Liu (Meta), Wenguang Mao (Meta), Dheevatsa Mudigere (NVIDIA), Maxim Naumov (Meta), Jongsoo Park (Meta), Mike Rabbat (Meta), Kaushik Rangadurai (Meta), Dennis van der Staay (Meta), Fei Tian (Meta), Rohan Varma (Meta), Sanjay Vishwakarma (Meta), Xunnan (Shawn) Xu (Meta), Jiyan Yang (Meta), Chunxing Yin (Meta), Gavin Zhang (Meta), Haoran Zhang (Meta), Haoyu Zhang (Meta), Chuanhao Zhuge (Meta), and Will Zou (Meta).
 
 ## 🏆 Competition Winner 🏆
 
@@ -74,9 +74,13 @@ A few notes on hyperparameters:
 
 - Here, `betas` refer to the hyperparameters used for the exponential moving average of the gradients and Shampoo preconditioners, while `grafting_beta2` corresponds to the `beta2` used specifically for exponential moving averaging of the grafted method. This is similar for `epsilon` and `grafting_epsilon`. As a first choice, we recommend setting `betas` equal to the previous `betas` and additionally setting `grafting_beta2` equal to `betas[1]`, and set `epsilon = 1e-12` and `grafting_epsilon` equal to the previous `epsilon`.
 
-- We also distinguish between `beta1` and iterate averaging. `beta1` (via `betas[0]`) corresponds to the EMA of the gradients (or gradient filtering), while iterate averaging (via `iterate_averaging_config`) provides momentum-like behavior through primal averaging. See [Example 7](#example-7-iterate-averaging-gpa-and-schedule-free) for details on configuring iterate averaging to achieve SGD momentum equivalence.
+- We also distinguish between `beta1` and iterate averaging. `beta1` (via `betas[0]`) corresponds to the EMA of the gradients (or gradient filtering), while iterate averaging (via `iterate_averaging_config`) provides momentum-like behavior. The supported configs are `ClassicMomentumConfig` (classic SGD-style momentum with dampening and Nesterov), `GeneralizedPrimalAveragingConfig` (GPA / primal averaging), and `ScheduleFreeConfig` (Schedule-Free). See [Example 7](#example-7-iterate-averaging-classic-momentum-gpa-and-schedule-free) for details and trade-offs.
 
-- We allow for decoupled and coupled weight decay. If one sets `use_decoupled_weight_decay=True`, then you are enabling AdamW-style weight decay, while `use_decoupled_weight_decay=False` corresponds to the normal L2-regularization style weight decay.
+- We allow for multiple weight decay strategies via `weight_decay_type`. The available options are:
+  - `WeightDecayType.L2`: Applies weight decay by adding a multiple of the weights to the gradient before preconditioning (L2 regularization).
+  - `WeightDecayType.DECOUPLED`: Applies AdamW-style decoupled weight decay independent of the preconditioned gradient.
+  - `WeightDecayType.CORRECTED`: Applies weight decay scaled by the learning rate divided by the maximum learning rate, independent of the preconditioned gradient.
+  - `WeightDecayType.INDEPENDENT`: Applies weight decay divided by the maximum learning rate, independent of the preconditioned gradient.
 
 - When setting `preconditioner_config` as an instance of `EigenvalueCorrectedShampooPreconditionerConfig` (see Example 5), there is typically no need to use learning rate grafting from Adam (`grafting_config=None`) and, when they are available, Adam's optimal `lr`, `betas`, and `weight_decay` should be a good starting point for further tuning. However, the case of `beta2=1.0`, i.e. an AdaGrad-like accumulation, has not been explored yet.  Also, in settings where Shampoo would usually graft its learning rate from SGD, grafting might still be beneficial.
 
@@ -100,8 +104,8 @@ we would instead use:
 ```python
 import torch
 from distributed_shampoo import (
+    ClassicMomentumConfig,
     DistributedShampoo,
-    GeneralizedPrimalAveragingConfig,
     SGDPreconditionerConfig,
 )
 
@@ -109,16 +113,17 @@ model = instantiate_model()
 
 optimizer = DistributedShampoo(
     model.parameters(),
-    lr=0.1,  # = 0.01 / (1 - 0.9) to account for primal averaging formulation
+    lr=0.01,
     betas=(0., 0.999),
     epsilon=1e-12,
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
     grafting_config=SGDPreconditionerConfig(),
-    iterate_averaging_config=GeneralizedPrimalAveragingConfig(
-        eval_interp_coeff=0.9,   # = momentum
-        train_interp_coeff=1.0,  # 1.0 for heavy-ball momentum
+    iterate_averaging_config=ClassicMomentumConfig(
+        momentum=0.9,
+        dampening=0.0,
+        use_nesterov=False,
     ),
 )
 ```
@@ -144,7 +149,7 @@ optimizer = Adam(
 we would instead use:
 ```python
 import torch
-from distributed_shampoo import AdamPreconditionerConfig, DistributedShampoo
+from distributed_shampoo import AdamPreconditionerConfig, DistributedShampoo, WeightDecayType
 
 model = instantiate_model()
 
@@ -156,7 +161,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=False,
+    weight_decay_type=WeightDecayType.L2,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-08,
@@ -183,7 +188,7 @@ optimizer = Adagrad(
 we would instead use:
 ```python
 import torch
-from distributed_shampoo import AdaGradPreconditionerConfig, DistributedShampoo
+from distributed_shampoo import AdaGradPreconditionerConfig, DistributedShampoo, WeightDecayType
 
 model = instantiate_model()
 
@@ -195,7 +200,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=False,
+    weight_decay_type=WeightDecayType.L2,
     grafting_config=AdaGradPreconditionerConfig(
         epsilon=1e-10,
     ),
@@ -222,7 +227,7 @@ optimizer = AdamW(
 we would instead use:
 ```python
 import torch
-from distributed_shampoo import AdamPreconditionerConfig, DistributedShampoo
+from distributed_shampoo import AdamPreconditionerConfig, DistributedShampoo, WeightDecayType
 
 model = instantiate_model()
 
@@ -234,7 +239,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-08,
@@ -265,6 +270,7 @@ import torch
 from distributed_shampoo import (
     DistributedShampoo,
     DefaultEigenvalueCorrectedShampooConfig,
+    WeightDecayType,
 )
 
 model = instantiate_model()
@@ -277,7 +283,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     # This can also be set to `DefaultSOAPConfig` which uses QR decompositions, hence is
     # less expensive and might thereby allow for a smaller `precondition_frequency`.
     preconditioner_config=DefaultEigenvalueCorrectedShampooConfig,
@@ -294,6 +300,7 @@ from distributed_shampoo import (
     NewtonSchulzOrthogonalizationConfig,
     SingleDeviceDistributedConfig,
     SpectralDescentPreconditionerConfig,
+    WeightDecayType,
 )
 
 
@@ -331,15 +338,54 @@ optimizer = DistributedShampoo(
         },
     ],
     weight_decay=1e-05,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
 )
 ```
 
 `SpectralDescentPreconditionerConfig` can also be used to implement other variations of spectral descent.
 
-### Example 7: Iterate Averaging (GPA and Schedule-Free)
+### Example 7: Iterate Averaging (Classic Momentum, GPA, and Schedule-Free)
+
+Distributed Shampoo supports three iterate averaging methods via `iterate_averaging_config`: `ClassicMomentumConfig` (classic SGD-style momentum applied to the preconditioned search direction), `GeneralizedPrimalAveragingConfig` (GPA), and `ScheduleFreeConfig`. GPA and Schedule-Free provide alternatives to traditional momentum with improved theoretical properties, while `ClassicMomentumConfig` provides a drop-in equivalent to PyTorch's SGD momentum semantics.
+
+#### Classic Momentum
+
+`ClassicMomentumConfig` applies classic SGD-style momentum directly to Shampoo's preconditioned search direction `P`:
+
+```
+M ← momentum * M + (1 - dampening) * P
+P ← (1 - dampening) * P + momentum * M    if use_nesterov
+P ← M                                      otherwise
+W ← W - lr * P
+```
+
+Unlike GPA, it requires no learning-rate rescaling and supports `dampening` and `use_nesterov`. It also does not require switching between `train()` / `eval()` modes (unlike Schedule-Free).
+
+```python
+from distributed_shampoo import (
+    ClassicMomentumConfig,
+    DistributedShampoo,
+    SGDPreconditionerConfig,
+)
+
+model = instantiate_model()
 
-Distributed Shampoo supports iterate averaging methods including Generalized Primal Averaging (GPA) and Schedule-Free optimization. These methods provide an alternative to traditional momentum with improved theoretical properties.
+# Shampoo with classic SGD momentum=0.9 (heavy-ball)
+optimizer = DistributedShampoo(
+    model.parameters(),
+    lr=0.01,
+    betas=(0.0, 1.0),
+    epsilon=1e-12,
+    max_preconditioner_dim=8192,
+    precondition_frequency=100,
+    preconditioner_config=SGDPreconditionerConfig(),
+    iterate_averaging_config=ClassicMomentumConfig(
+        momentum=0.9,
+        dampening=0.0,
+        use_nesterov=False,
+    ),
+)
+```
 
 #### Generalized Primal Averaging (GPA)
 
@@ -408,15 +454,15 @@ optimizer.eval()
 
 #### Migration from Previous Momentum Parameters
 
-The previous `momentum`, `dampening`, and `use_nesterov` parameters have been replaced by iterate averaging configs. Here are the equivalences:
+The previous `momentum`, `dampening`, and `use_nesterov` parameters have been replaced by iterate averaging configs. The simplest drop-in replacement is `ClassicMomentumConfig`, which preserves SGD's exact momentum/dampening/Nesterov semantics with no learning-rate rescaling. `GeneralizedPrimalAveragingConfig` is an alternative formulation with different theoretical properties (and requires `lr = lr / (1 - β)` for heavy-ball equivalence).
 
-| Previous Configuration | New Configuration |
-|------------------------|-------------------|
-| `momentum=β, dampening=0, use_nesterov=False` | `GeneralizedPrimalAveragingConfig(eval_interp_coeff=β, train_interp_coeff=1.0)` with `lr = lr / (1-β)` |
-| `momentum=β, dampening=0, use_nesterov=True` | `GeneralizedPrimalAveragingConfig(eval_interp_coeff=β, train_interp_coeff=β)` with `lr = lr / (1-β)` |
-| `momentum=β, dampening=d` (d≠0) | No direct equivalent. Dampening is not supported in iterate averaging. |
+| Previous Configuration | Recommended (`ClassicMomentumConfig`) | Equivalent GPA Configuration |
+|------------------------|---------------------------------------|------------------------------|
+| `momentum=β, dampening=0, use_nesterov=False` | `ClassicMomentumConfig(momentum=β)` (no `lr` change) | `GeneralizedPrimalAveragingConfig(eval_interp_coeff=β, train_interp_coeff=1.0)` with `lr = lr / (1-β)` |
+| `momentum=β, dampening=0, use_nesterov=True` | `ClassicMomentumConfig(momentum=β, use_nesterov=True)` (no `lr` change) | `GeneralizedPrimalAveragingConfig(eval_interp_coeff=β, train_interp_coeff=β)` with `lr = lr / (1-β)` |
+| `momentum=β, dampening=d` (d≠0) | `ClassicMomentumConfig(momentum=β, dampening=d)` | No direct equivalent — dampening is not supported in GPA. |
 
-**Note on LaProp:** The previous momentum implementation (sometimes called LaProp) is mathematically equivalent to the heavy-ball/primal averaging formulation when `dampening=0`. Use the heavy-ball configuration above.
+**Note on LaProp:** The previous momentum implementation (sometimes called LaProp) is mathematically equivalent to either `ClassicMomentumConfig` (when `dampening=0`) or the heavy-ball/primal averaging formulation. Both are valid migrations.
 
 **Note on LaPropW:** The original LaPropW from the paper includes additional weight decay handling that may differ slightly from the iterate averaging formulation. For most practical purposes, the heavy-ball configuration provides equivalent behavior.
 
@@ -444,6 +490,7 @@ from distributed_shampoo import (
     AdamPreconditionerConfig,
     DDPDistributedConfig,
     DistributedShampoo,
+    WeightDecayType,
 )
 from torch import nn
 
@@ -473,7 +520,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-12,
@@ -506,6 +553,7 @@ from distributed_shampoo import (
     compile_fsdp_parameter_metadata,
     DistributedShampoo,
     FSDPDistributedConfig,
+    WeightDecayType,
 )
 
 LOCAL_RANK = int(os.environ["LOCAL_RANK"])
@@ -531,7 +579,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-12,
@@ -558,6 +606,7 @@ from distributed_shampoo import (
     compile_fsdp_parameter_metadata,
     DistributedShampoo,
     HSDPDistributedConfig,
+    WeightDecayType,
 )
 
 LOCAL_RANK = int(os.environ["LOCAL_RANK"])
@@ -590,7 +639,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-12,
@@ -620,6 +669,7 @@ from distributed_shampoo import (
     AdamPreconditionerConfig,
     DistributedShampoo,
     FullyShardDistributedConfig,
+    WeightDecayType,
 )
 
 LOCAL_RANK = int(os.environ["LOCAL_RANK"])
@@ -645,7 +695,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-12,
@@ -671,6 +721,7 @@ from distributed_shampoo import (
     AdamPreconditionerConfig,
     DistributedShampoo,
     HybridShardDistributedConfig,
+    WeightDecayType,
 )
 
 LOCAL_RANK = int(os.environ["LOCAL_RANK"])
@@ -703,7 +754,7 @@ optimizer = DistributedShampoo(
     weight_decay=1e-05,
     max_preconditioner_dim=8192,
     precondition_frequency=100,
-    use_decoupled_weight_decay=True,
+    weight_decay_type=WeightDecayType.DECOUPLED,
     grafting_config=AdamPreconditionerConfig(
         beta2=0.999,
         epsilon=1e-12,
diff --git a/distributed_shampoo/distributed_shampoo.py b/distributed_shampoo/distributed_shampoo.py
index 96288f1..5325a30 100644
--- a/distributed_shampoo/distributed_shampoo.py
+++ b/distributed_shampoo/distributed_shampoo.py
@@ -11,6 +11,7 @@
 import math
 import operator
 from collections.abc import Callable
+from contextlib import nullcontext
 from dataclasses import asdict, fields, is_dataclass
 from types import LambdaType
 from typing import Any, overload
@@ -23,6 +24,7 @@
     HybridShardLosslessDistributor,
 )
 from distributed_shampoo.distributor.shampoo_ddp_distributor import DDPDistributor
+from distributed_shampoo.distributor.shampoo_dist_utils import cuda_stream_context
 from distributed_shampoo.distributor.shampoo_distributor import (
     Distributor,
     DistributorInterface,
@@ -65,7 +67,11 @@
     BaseShampooPreconditionerConfig,
     BETA3,
     BETAS,
+    ClassicMomentumConfig,
+    ConcurrencyConfig,
+    CUDA_STREAM,
     DDPDistributedConfig,
+    DefaultConcurrencyConfig,
     DefaultShampooConfig,
     DefaultShampooRuntimeConfig,
     DefaultSingleDeviceDistributedConfig,
@@ -76,6 +82,7 @@
     EigendecomposedShampooPreconditionerConfig,
     EigenvalueCorrectedShampooPreconditionerConfig,
     EPSILON,
+    EVAL_INTERP_COEFF,
     FILTERED_GRAD,
     FILTERED_GRAD_LIST,
     FSDPDistributedConfig,
@@ -95,8 +102,11 @@
     MASKED_BLOCKED_GRADS,
     MASKED_BLOCKED_PARAMS,
     MASKED_FILTERED_GRAD_LIST,
+    MASKED_MOMENTUM_BUFFER_LIST,
     MASKED_WEIGHT_BUFFER_LIST,
     MAX_PRECONDITIONER_DIM,
+    MOMENTUM_BUFFER,
+    MOMENTUM_BUFFER_LIST,
     PARAMS,
     PEAK_LR,
     PRECONDITION_FREQUENCY,
@@ -116,6 +126,7 @@
     SpectralDescentPreconditionerConfig,
     START_PRECONDITIONING_STEP,
     STEP,
+    TRAIN_INTERP_COEFF,
     TRAIN_MODE,
     USE_BIAS_CORRECTION,
     USE_PIN_MEMORY,
@@ -129,7 +140,7 @@
     extract_state_dict_content,
     update_param_state_dict_object,
 )
-from distributed_shampoo.utils.shampoo_utils import compress_list
+from distributed_shampoo.utils.shampoo_utils import compress_list, split_param_groups
 from torch.optim.optimizer import Optimizer, ParamsT, StateDict
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -347,7 +358,7 @@ class DistributedShampoo(torch.optim.Optimizer):
 
     """
 
-    def __init__(
+    def __init__(  # noqa: C901
         self,
         params: ParamsT,
         lr: float = 1e-2,
@@ -367,6 +378,7 @@ def __init__(
         distributed_config: DistributedConfig = DefaultSingleDeviceDistributedConfig,
         preconditioner_config: PreconditionerConfig = DefaultShampooConfig,
         shampoo_runtime_config: ShampooRuntimeConfig = DefaultShampooRuntimeConfig,
+        concurrency_config: ConcurrencyConfig = DefaultConcurrencyConfig,
     ) -> None:
         super().__init__(
             params,
@@ -393,6 +405,9 @@ def __init__(
         self.register_load_state_dict_pre_hook(self._pre_load_state_dict_hook)
         self.register_load_state_dict_post_hook(self._post_load_state_dict_hook)
 
+        # Split FSDP/HSDP param groups with num_sub_groups > 1 for multi-stream overlap.
+        self.param_groups[:] = split_param_groups(self.param_groups)
+
         def param_group_hyperparameter_check(param_group: dict[str, Any]) -> None:
             if not param_group[LR] >= 0.0:
                 raise ValueError(f"Invalid {param_group[LR]=}. Must be >= 0.0.")
@@ -524,6 +539,7 @@ def param_group_hyperparameter_check(param_group: dict[str, Any]) -> None:
             shampoo_pt2_compile_config
         )
         self._runtime_config: ShampooRuntimeConfig = shampoo_runtime_config
+        self._concurrency_config: ConcurrencyConfig = concurrency_config
 
         # Initialize list containing group state dictionaries.
         self._per_group_state_lists: list[dict[str, Any]] = [
@@ -552,45 +568,7 @@ def _instantiate_distributor(self) -> None:
             if not group[PARAMS]:
                 raise ValueError(f"Shampoo got an empty parameter {group=}")
 
-            match group[DISTRIBUTED_CONFIG]:
-                case SingleDeviceDistributedConfig():
-                    distributor_cls: type[DistributorInterface] = Distributor
-                case HSDPDistributedConfig():
-                    distributor_cls = HSDPDistributor
-                case HybridShardDistributedConfig(
-                    param_assignment_strategy=FSDPParamAssignmentStrategy.DEFAULT
-                ):
-                    distributor_cls = HybridShardDistributor
-                case (
-                    HybridShardDistributedConfig(
-                        param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE
-                    )
-                    | HybridShardDistributedConfig(
-                        param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN
-                    )
-                ):
-                    distributor_cls = HybridShardLosslessDistributor
-                case DDPDistributedConfig():
-                    distributor_cls = DDPDistributor
-                case FSDPDistributedConfig():
-                    distributor_cls = FSDPDistributor
-                case FullyShardDistributedConfig(
-                    param_assignment_strategy=FSDPParamAssignmentStrategy.DEFAULT
-                ):
-                    distributor_cls = FullyShardDistributor
-                case (
-                    FullyShardDistributedConfig(
-                        param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE
-                    )
-                    | FullyShardDistributedConfig(
-                        param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN
-                    )
-                ):
-                    distributor_cls = FullyShardLosslessDistributor
-                case _:
-                    raise NotImplementedError(
-                        f"{group[DISTRIBUTED_CONFIG]=} not supported!"
-                    )
+            distributor_cls = self._select_distributor_class(group[DISTRIBUTED_CONFIG])
 
             # Instantiate distributors for each group.
             state_lists[DISTRIBUTOR] = distributor_cls(group, self._runtime_config)
@@ -614,6 +592,57 @@ def _instantiate_distributor(self) -> None:
             # First PREVIOUS_GRAD_SELECTOR is set to None.
             state_lists[PREVIOUS_GRAD_SELECTOR] = None
 
+    def _select_distributor_class(
+        self, distributed_config: DistributedConfig
+    ) -> type[DistributorInterface]:
+        """Map a per-group ``DistributedConfig`` to its ``DistributorInterface``.
+
+        Factored out of ``_instantiate_distributor`` so the config-to-distributor
+        dispatch has a single override point. Subclasses can inject a custom
+        distributor (e.g. expert-parallel) by handling their own config types
+        here and delegating everything else to
+        ``super()._select_distributor_class(...)`` -- without reimplementing the
+        per-group instantiation loop or its bookkeeping.
+        """
+        match distributed_config:
+            case SingleDeviceDistributedConfig():
+                distributor_cls: type[DistributorInterface] = Distributor
+            case HSDPDistributedConfig():
+                distributor_cls = HSDPDistributor
+            case HybridShardDistributedConfig(
+                param_assignment_strategy=FSDPParamAssignmentStrategy.DEFAULT
+            ):
+                distributor_cls = HybridShardDistributor
+            case (
+                HybridShardDistributedConfig(
+                    param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE
+                )
+                | HybridShardDistributedConfig(
+                    param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN
+                )
+            ):
+                distributor_cls = HybridShardLosslessDistributor
+            case DDPDistributedConfig():
+                distributor_cls = DDPDistributor
+            case FSDPDistributedConfig():
+                distributor_cls = FSDPDistributor
+            case FullyShardDistributedConfig(
+                param_assignment_strategy=FSDPParamAssignmentStrategy.DEFAULT
+            ):
+                distributor_cls = FullyShardDistributor
+            case (
+                FullyShardDistributedConfig(
+                    param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE
+                )
+                | FullyShardDistributedConfig(
+                    param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN
+                )
+            ):
+                distributor_cls = FullyShardLosslessDistributor
+            case _:
+                raise NotImplementedError(f"{distributed_config=} not supported!")
+        return distributor_cls
+
     @torch.no_grad()
     def _initialize_blocked_parameters_state(self) -> None:
         for state_lists in self._per_group_state_lists:
@@ -785,6 +814,18 @@ def _instantiate_lr_tensors(self) -> None:
             state_lists[LR_CPU_PINNED] = torch.empty(
                 (), dtype=torch.float, pin_memory=group[USE_PIN_MEMORY]
             )
+            # Pre-allocate persistent interp coeff tensors for iterate averaging.
+            # Using persistent tensors avoids Dynamo recompilation: Dynamo treats
+            # tensor inputs as dynamic (doesn't specialize on value), unlike Python
+            # floats which trigger recompilation on every value change. This is
+            # critical for ScheduleFreeConfig where eval_interp_coeff changes every step.
+            # For ClassicMomentumConfig, these remain at 0.0 (no interpolation needed).
+            state_lists[TRAIN_INTERP_COEFF] = torch.zeros(
+                (), dtype=torch.float, device=device
+            )
+            state_lists[EVAL_INTERP_COEFF] = torch.zeros(
+                (), dtype=torch.float, device=device
+            )
 
     @torch.no_grad()
     def _instantiate_filtered_grads(self) -> None:
@@ -827,68 +868,92 @@ def _instantiate_filtered_grads(self) -> None:
 
     @torch.no_grad()
     def _instantiate_iterate_averaging(self) -> None:
-        # NOTE: Since we are using the memory-efficient implementation of GPA and Schedule-Free, when iterate
-        # averaging is enabled, we instantiate a weight buffer (block_state[WEIGHT_BUFFER]) that stores each
-        # parameter's "z" sequence. The current parameters are instead treated as the "y" sequence where the
-        # gradient is computed.
-        #
-        # To use the "x" sequence, one must enable the eval mode for the optimizer, which switches the "y"
-        # sequence to the "x" sequence. Train mode switches from the "x" sequence to the "y" sequence.
         for state_lists, group in zip(
             self._per_group_state_lists, self.param_groups, strict=True
         ):
-            if group[ITERATE_AVERAGING_CONFIG] is None:
-                continue
+            match group[ITERATE_AVERAGING_CONFIG]:
+                case None:
+                    continue
+                case ClassicMomentumConfig():
+                    self._instantiate_classical_momentum(state_lists, group)
+                case GeneralizedPrimalAveragingConfig() | ScheduleFreeConfig():
+                    self._instantiate_gpa_schedule_free(state_lists, group)
+                case _:
+                    raise NotImplementedError(
+                        f"{group[ITERATE_AVERAGING_CONFIG]=} is not supported! "
+                        "Supported configs: ClassicMomentumConfig, GeneralizedPrimalAveragingConfig, ScheduleFreeConfig."
+                    )
 
-            # Construct local weight buffer list.
-            local_weight_buffer_list = []
-            for block, block_info in zip(
-                state_lists[DISTRIBUTOR].local_blocked_params,
-                state_lists[DISTRIBUTOR].local_block_info_list,
-                strict=True,
-            ):
-                assert (
-                    block_index := block_info.composable_block_ids[1]
-                ) in self.state[block_info.param], (
-                    f"{block_index=} not found in {self.state[block_info.param]=}. "
-                    "Please check the initialization of self.state[block_info.param][block_index] "
-                    "within _initialize_blocked_parameters_state, and check the initialization of BlockInfo "
-                    "within Distributor for the correctness of block_index."
-                )
-                block_state = self.state[block_info.param][block_index]
+    def _instantiate_classical_momentum(
+        self, state_lists: dict[str, Any], group: dict[str, Any]
+    ) -> None:
+        self._allocate_block_buffer_list(
+            state_lists=state_lists,
+            buffer_key=MOMENTUM_BUFFER,
+            list_key=MOMENTUM_BUFFER_LIST,
+            masked_list_key=MASKED_MOMENTUM_BUFFER_LIST,
+        )
 
-                block_state[WEIGHT_BUFFER] = block_info.allocate_zeros_tensor(
-                    size=block.size(),
-                    dtype=block.dtype,
-                    device=block.device,
-                )
-                # Get the local tensor from the DTensor and copy into it.
-                local_weight_buffer = block_info.get_tensor(block_state[WEIGHT_BUFFER])
-                local_weight_buffer.copy_(block)
-                local_weight_buffer_list.append(local_weight_buffer)
+    def _instantiate_gpa_schedule_free(
+        self, state_lists: dict[str, Any], group: dict[str, Any]
+    ) -> None:
+        # The "z" sequence is initialized from the current parameters (the "y" sequence
+        # where gradients are computed). Train/eval mode switches between "y" and the
+        # averaged "x" sequence.
+        self._allocate_block_buffer_list(
+            state_lists=state_lists,
+            buffer_key=WEIGHT_BUFFER,
+            list_key=WEIGHT_BUFFER_LIST,
+            masked_list_key=MASKED_WEIGHT_BUFFER_LIST,
+            init_from_param=True,
+        )
 
-            state_lists[WEIGHT_BUFFER_LIST] = local_weight_buffer_list
-            # Here, we set masked weight buffer list to weight buffer list because we assume
-            # all parameters are active.
-            state_lists[MASKED_WEIGHT_BUFFER_LIST] = state_lists[WEIGHT_BUFFER_LIST]
-
-            # Instantiate summed learning rate tensor used by iterate averaging
-            # to track the cumulative learning rate across steps.
-            state_lists[LR_SUM] = torch.tensor(
-                0.0,
-                dtype=torch.float,
-                device=group[PARAMS][0].device,
+        state_lists[LR_SUM] = torch.tensor(
+            0.0, dtype=torch.float, device=group[PARAMS][0].device
+        )
+        state_lists[TRAIN_MODE] = torch.tensor(True, dtype=torch.bool, device="cpu")
+
+        # Stored under the first parameter's state so train mode and LR sum are checkpointed.
+        self.state[group[PARAMS][0]][TRAIN_MODE] = state_lists[TRAIN_MODE]
+        self.state[group[PARAMS][0]][LR_SUM] = state_lists[LR_SUM]
+
+    def _allocate_block_buffer_list(
+        self,
+        state_lists: dict[str, Any],
+        buffer_key: str,
+        list_key: str,
+        masked_list_key: str,
+        init_from_param: bool = False,
+    ) -> None:
+        local_buffer_list = []
+        for block, block_info in zip(
+            state_lists[DISTRIBUTOR].local_blocked_params,
+            state_lists[DISTRIBUTOR].local_block_info_list,
+            strict=True,
+        ):
+            assert (block_index := block_info.composable_block_ids[1]) in self.state[
+                block_info.param
+            ], (
+                f"{block_index=} not found in {self.state[block_info.param]=}. "
+                "Please check the initialization of self.state[block_info.param][block_index] "
+                "within _initialize_blocked_parameters_state, and check the initialization of BlockInfo "
+                "within Distributor for the correctness of block_index."
             )
+            block_state = self.state[block_info.param][block_index]
 
-            # Instantiate a single boolean tensor on CPU for each group in order
-            # to track the train and eval mode of each parameter group in the optimizer.
-            # This is used by iterate averaging to determine whether to apply train or eval updates.
-            state_lists[TRAIN_MODE] = torch.tensor(True, dtype=torch.bool, device="cpu")
+            block_state[buffer_key] = block_info.allocate_zeros_tensor(
+                size=block.size(),
+                dtype=block.dtype,
+                device=block.device,
+            )
+            local_buffer = block_info.get_tensor(block_state[buffer_key])
+            if init_from_param:
+                local_buffer.copy_(block)
+            local_buffer_list.append(local_buffer)
 
-            # In order to ensure that the train mode and learning rate sum are checkpointed correctly,
-            # we store it as a tensor (which is replicated across all devices) under the first parameter's state.
-            self.state[group[PARAMS][0]][TRAIN_MODE] = state_lists[TRAIN_MODE]
-            self.state[group[PARAMS][0]][LR_SUM] = state_lists[LR_SUM]
+        state_lists[list_key] = local_buffer_list
+        # Masked list aliases the full list since all parameters are assumed active.
+        state_lists[masked_list_key] = state_lists[list_key]
 
     @torch.no_grad()
     def _instantiate_per_group_step(
@@ -897,7 +962,10 @@ def _instantiate_per_group_step(
         # Use PT2 to compile the step function for each parameter group.
         self._per_group_step: Callable[..., None] = (
             torch.compile(
-                self._per_group_step_impl, **asdict(shampoo_pt2_compile_config)
+                # pyrefly: ignore [bad-argument-type]
+                self._per_group_step_impl,
+                # pyrefly: ignore [bad-argument-type]
+                **asdict(shampoo_pt2_compile_config),
             )
             if shampoo_pt2_compile_config is not None
             else self._per_group_step_impl
@@ -907,6 +975,23 @@ def _instantiate_per_group_step(
                 f"DistributedShampoo optimizer initialization is using {shampoo_pt2_compile_config=}"
             )
 
+        # Per-group CUDA stream isolation requires CUDA params on every group.
+        all_cuda = all(g[PARAMS][0].device.type == "cuda" for g in self.param_groups)
+        self._use_cuda_streams: bool = (
+            self._concurrency_config.enable_cuda_stream_for_param_groups and all_cuda
+        )
+        if self._use_cuda_streams:
+            for i, (state_lists, group) in enumerate(
+                zip(self._per_group_state_lists, self.param_groups, strict=True)
+            ):
+                stream = torch.cuda.Stream(device=group[PARAMS][0].device)
+                state_lists[CUDA_STREAM] = stream
+                logger.info(
+                    "param_group_%d: stream_id=%s",
+                    i,
+                    stream.stream_id,
+                )
+
     @staticmethod
     @torch.no_grad()
     def _mask_state_lists(
@@ -963,11 +1048,16 @@ def _mask_state_lists(
                 state_lists[FILTERED_GRAD_LIST],
                 state_lists[DISTRIBUTOR].local_grad_selector,
             )
-        if group[ITERATE_AVERAGING_CONFIG] is not None:
+        if WEIGHT_BUFFER_LIST in state_lists:
             state_lists[MASKED_WEIGHT_BUFFER_LIST] = compress_list(
                 state_lists[WEIGHT_BUFFER_LIST],
                 state_lists[DISTRIBUTOR].local_grad_selector,
             )
+        if MOMENTUM_BUFFER_LIST in state_lists:
+            state_lists[MASKED_MOMENTUM_BUFFER_LIST] = compress_list(
+                state_lists[MOMENTUM_BUFFER_LIST],
+                state_lists[DISTRIBUTOR].local_grad_selector,
+            )
 
     @torch.no_grad()
     @torch.compiler.disable
@@ -1090,7 +1180,7 @@ def _compute_filtered_grad_list(
         return masked_filtered_grad_list
 
     @torch.no_grad()
-    def _apply_decoupled_or_corrected_weight_decay(
+    def _apply_weight_decay(
         self,
         state_lists: dict[str, Any],
         masked_blocked_search_directions: tuple[torch.Tensor, ...],
@@ -1098,45 +1188,108 @@ def _apply_decoupled_or_corrected_weight_decay(
         weight_decay: float,
         peak_lr: float,
         weight_decay_type: WeightDecayType,
+        iterate_averaging_enabled: bool,
     ) -> None:
-        if weight_decay != 0.0 and weight_decay_type in (
-            WeightDecayType.DECOUPLED,
-            WeightDecayType.CORRECTED,
-            WeightDecayType.INDEPENDENT,
-        ):
-            match weight_decay_type:
-                case WeightDecayType.DECOUPLED:
-                    alpha = weight_decay
-                case WeightDecayType.CORRECTED:
-                    alpha = weight_decay * lr.item() / peak_lr
-                case WeightDecayType.INDEPENDENT:
-                    alpha = weight_decay / peak_lr
-                case _:
-                    raise ValueError(
-                        f"Invalid weight decay type: {weight_decay_type=}!"
-                    )
+        """Apply weight decay (decoupled, corrected, or independent).
+
+        Two cases require different strategies:
+
+        Without iterate averaging: adds weight decay to the search direction
+        (P += alpha * W). This is mathematically equivalent to multiplicative
+        shrinkage (W *= 1 - lr * alpha), but the additive form is required
+        because MASKED_BLOCKED_PARAMS only contains the local subset of
+        blocks assigned to each rank. In-place shrinkage would only decay
+        local blocks, causing params to diverge across ranks. By folding
+        weight decay into the search direction, it flows through the
+        all-gather in update_params and is applied uniformly.
+
+        With iterate averaging: shrinks Z directly (Z *= 1 - lr * alpha).
+        Weight decay cannot be folded into the search direction P here
+        because P is reused by the primal averaging update; the weight
+        decay term would leak into the W update via the primal averaging
+        coefficients. Shrinking Z directly is safe because Z is local
+        optimizer state (not shared across ranks via the distributor).
+        """
+        if weight_decay == 0.0 or weight_decay_type == WeightDecayType.L2:
+            return
+
+        match weight_decay_type:
+            case WeightDecayType.DECOUPLED:
+                alpha = weight_decay
+            case WeightDecayType.CORRECTED:
+                # Use LR_CPU_PINNED instead of lr (GPU tensor) to avoid
+                # a host-device sync from lr.item().
+                alpha = weight_decay * state_lists[LR_CPU_PINNED].item() / peak_lr
+            case WeightDecayType.INDEPENDENT:
+                alpha = weight_decay / peak_lr
+            case _:
+                raise ValueError(f"Invalid weight decay type: {weight_decay_type=}!")
+
+        if iterate_averaging_enabled:
+            torch._foreach_mul_(state_lists[MASKED_WEIGHT_BUFFER_LIST], 1 - lr * alpha)
+        else:
             torch._foreach_add_(
                 masked_blocked_search_directions,
                 state_lists[MASKED_BLOCKED_PARAMS],
                 alpha=alpha,
             )
 
+    @torch.no_grad()
+    def _apply_classic_momentum(
+        self,
+        state_lists: dict[str, Any],
+        masked_blocked_search_directions: tuple[torch.Tensor, ...],
+        momentum_param: float,
+        dampening: float,
+        use_nesterov: bool,
+    ) -> None:
+        """Apply classic SGD-style momentum to the search directions.
+
+        Updates the momentum buffer and modifies search directions in-place:
+            M <- momentum * M + (1 - dampening) * P
+            P <- (1 - dampening) * P + momentum * M     if use_nesterov
+            P <- M                                       otherwise
+        """
+        torch._foreach_mul_(state_lists[MASKED_MOMENTUM_BUFFER_LIST], momentum_param)
+        torch._foreach_add_(
+            state_lists[MASKED_MOMENTUM_BUFFER_LIST],
+            masked_blocked_search_directions,
+            alpha=1 - dampening,
+        )
+
+        if use_nesterov:
+            torch._foreach_mul_(
+                masked_blocked_search_directions,
+                1 - dampening,
+            )
+            torch._foreach_add_(
+                masked_blocked_search_directions,
+                state_lists[MASKED_MOMENTUM_BUFFER_LIST],
+                alpha=momentum_param,
+            )
+        else:
+            torch._foreach_copy_(
+                masked_blocked_search_directions,
+                state_lists[MASKED_MOMENTUM_BUFFER_LIST],
+            )
+
     @staticmethod
     @torch.no_grad()
     def _get_train_and_eval_interp_coeffs(
-        iterate_averaging_config: IterateAveragingConfig | None,
+        iterate_averaging_config: GeneralizedPrimalAveragingConfig | ScheduleFreeConfig,
         lr: torch.Tensor | None,
         state_lists: dict[str, Any],
-    ) -> tuple[float, float]:
-        if iterate_averaging_config is None:
-            return 0.0, 0.0
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Persistent tensors pre-allocated in _instantiate_lr_tensors; filled in place below.
+        train_interp_coeff = state_lists[TRAIN_INTERP_COEFF]
+        eval_interp_coeff = state_lists[EVAL_INTERP_COEFF]
 
         match iterate_averaging_config:
             case GeneralizedPrimalAveragingConfig():
-                train_interp_coeff = iterate_averaging_config.train_interp_coeff
-                eval_interp_coeff = iterate_averaging_config.eval_interp_coeff
+                train_interp_coeff.fill_(iterate_averaging_config.train_interp_coeff)
+                eval_interp_coeff.fill_(iterate_averaging_config.eval_interp_coeff)
             case ScheduleFreeConfig():
-                train_interp_coeff = iterate_averaging_config.train_interp_coeff
+                train_interp_coeff.fill_(iterate_averaging_config.train_interp_coeff)
                 if lr is not None:
                     # Based on Equation (23) in The Road Less Scheduled (https://arxiv.org/pdf/2405.15682).
                     # In the simplest case, computes eval_interp_coeff = 1 - gamma_t^2 / \sum_{i=1}^t gamma_i^2.
@@ -1150,15 +1303,19 @@ def _get_train_and_eval_interp_coeffs(
                         lr, iterate_averaging_config.eval_coeff_lr_power
                     )
                     lr_sum += lr_power
-                    eval_interp_coeff = 1.0 - (lr_power / lr_sum).item()
+                    # Keep as tensor — avoids .item() which creates a Python float
+                    # that Dynamo would specialize on, causing recompilation every step.
+                    eval_interp_coeff.copy_(1.0 - lr_power / lr_sum)
                 else:
                     logger.warning(
                         "lr or lr_sum not provided to _get_train_and_eval_interp_coeffs; returning eval_interp_coeff = 0.0!"
                     )
-                    eval_interp_coeff = 0.0
+                    eval_interp_coeff.fill_(0.0)
             case _:
                 raise ValueError(
-                    f"Unsupported iterate averaging config {iterate_averaging_config=}!"
+                    f"Iterate averaging config {iterate_averaging_config=} is not supported "
+                    "for retrieving train and eval interpolation coefficients; only "
+                    "GeneralizedPrimalAveragingConfig and ScheduleFreeConfig are supported."
                 )
 
         return (train_interp_coeff, eval_interp_coeff)
@@ -1168,14 +1325,9 @@ def _apply_in_place_primal_averaging(
         self,
         state_lists: dict[str, Any],
         masked_blocked_search_directions: tuple[torch.Tensor, ...],
-        train_interp_coeff: float,
-        eval_interp_coeff: float,
+        train_interp_coeff: torch.Tensor,
+        eval_interp_coeff: torch.Tensor,
     ) -> None:
-        # If train_interp_coeff = 0, then GPA is equivalent to the base optimizer, so we
-        # return immediately.
-        if train_interp_coeff == 0.0:
-            return
-
         # Raise user error when attempting to perform primal averaging but not in train mode.
         if not state_lists[TRAIN_MODE]:
             raise RuntimeError(
@@ -1199,13 +1351,14 @@ 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),
+            (1 - train_interp_coeff * eval_interp_coeff),  # type: ignore
         )
         # This computes: (1 - mu_x) * (Z_old - Y) - (1 - mu_x * mu_y) * lr * P.
+        # pyrefly: ignore [no-matching-overload]
         torch._foreach_add_(
             masked_blocked_search_directions,
             z_minus_y_term,
-            alpha=1 - eval_interp_coeff,
+            alpha=(1 - eval_interp_coeff),  # type: ignore
         )
 
     @torch.no_grad()
@@ -1223,8 +1376,13 @@ def _compute_search_directions(
         perform_amortized_computation: bool,
         use_bias_correction: bool,
         use_grafting_method: bool,
-        train_interp_coeff: float,
-        eval_interp_coeff: float,
+        iterate_averaging_enabled: bool,
+        classic_momentum_enabled: bool,
+        train_interp_coeff: torch.Tensor,
+        eval_interp_coeff: torch.Tensor,
+        momentum_param: float,
+        dampening: float,
+        use_nesterov: bool,
     ) -> tuple[torch.Tensor, ...]:
         # Incorporate L2-regularization or (coupled) weight decay if enabled.
         #   G <- G + weight_decay * W
@@ -1279,34 +1437,47 @@ def _compute_search_directions(
             grafting_config_not_none,
         )
 
-        # Incorporate decoupled or corrected weight decay into search direction if enabled.
-        #   P <- P + weight_decay * W                       if decoupled
-        #   P <- P + weight_decay * (lr / peak_lr) * W      if corrected
-        #   P <- P + weight_decay * (1 / peak_lr) * W       if independent
-        self._apply_decoupled_or_corrected_weight_decay(
+        # Apply weight decay (decoupled, corrected, or independent).
+        self._apply_weight_decay(
             state_lists,
             masked_blocked_search_directions,
             lr,
             weight_decay,
             peak_lr,
             weight_decay_type,
+            iterate_averaging_enabled,
         )
 
+        # Apply classic momentum if enabled (ClassicMomentumConfig).
+        #   M <- momentum * M + (1 - dampening) * P
+        #   P <- (1 - dampening) * P + momentum * M     if use_nesterov
+        #   P <- M                                       otherwise
+        if classic_momentum_enabled:
+            self._apply_classic_momentum(
+                state_lists,
+                masked_blocked_search_directions,
+                momentum_param,
+                dampening,
+                use_nesterov,
+            )
+
         # Multiplies the learning rate to the search direction / update.
         torch._foreach_mul_(masked_blocked_search_directions, -lr)
 
-        # Incorporates primal averaging into the search direction if enabled.
+        # Incorporates primal averaging into the search direction if enabled
+        # (GeneralizedPrimalAveragingConfig / ScheduleFreeConfig).
         # NOTE: When primal averaging is enabled, we set Y = W in train mode.
         #   P <- (1 - mu_x * mu_y) * P + (1 - mu_x) * (Z - W)
         #
         # This is equivalent to the expanded update:
         #   P <- mu_x * Y + (1 - mu_x) * Z - (1 - mu_x * mu_y) * lr * P
-        self._apply_in_place_primal_averaging(
-            state_lists,
-            masked_blocked_search_directions,
-            train_interp_coeff,
-            eval_interp_coeff,
-        )
+        if iterate_averaging_enabled:
+            self._apply_in_place_primal_averaging(
+                state_lists,
+                masked_blocked_search_directions,
+                train_interp_coeff,
+                eval_interp_coeff,
+            )
 
         return masked_blocked_search_directions
 
@@ -1325,8 +1496,13 @@ def _per_group_step_impl(
         perform_amortized_computation: bool,
         use_bias_correction: bool,
         use_grafting_method: bool,
-        train_interp_coeff: float,
-        eval_interp_coeff: float,
+        iterate_averaging_enabled: bool,
+        classic_momentum_enabled: bool,
+        train_interp_coeff: torch.Tensor,
+        eval_interp_coeff: torch.Tensor,
+        momentum_param: float,
+        dampening: float,
+        use_nesterov: bool,
     ) -> None:
         # This method computes search directions and updates parameters in one step
         # It's designed to be compiled with PyTorch 2.0 for performance optimization
@@ -1349,8 +1525,13 @@ def _per_group_step_impl(
                 perform_amortized_computation=perform_amortized_computation,
                 use_bias_correction=use_bias_correction,
                 use_grafting_method=use_grafting_method,
+                iterate_averaging_enabled=iterate_averaging_enabled,
+                classic_momentum_enabled=classic_momentum_enabled,
                 train_interp_coeff=train_interp_coeff,
                 eval_interp_coeff=eval_interp_coeff,
+                momentum_param=momentum_param,
+                dampening=dampening,
+                use_nesterov=use_nesterov,
             )
             # Only update parameters if there are gradients to use
             # Otherwise, return an empty tuple to avoid unnecessary computation
@@ -1358,6 +1539,114 @@ def _per_group_step_impl(
             else ()
         )
 
+    @torch.no_grad()
+    def _group_step_body(
+        self,
+        state_lists: dict[str, Any],
+        group: dict[str, Any],
+        per_group_step: Callable[..., None],
+    ) -> None:
+        """Core per-group step logic; called from step() once per param group,
+        wrapped in cuda_stream_context (or nullcontext) by the caller."""
+        # Construct blocked gradient list.
+        state_lists[MASKED_BLOCKED_GRADS] = state_lists[
+            DISTRIBUTOR
+        ].merge_and_block_gradients()
+
+        # Based on the current block selector, mask lists of parameters and optimizer states.
+        DistributedShampoo._mask_state_lists(
+            state_lists=state_lists,
+            group=group,
+            shampoo_pt2_enabled=self._shampoo_pt2_compile_config is not None,
+        )
+
+        # Iterate group step counter and define Python scalar step.
+        step = state_lists[STEP].add_(1)
+        step_val = step.item()
+        # NOTE: Reuse pre-allocated lr tensors to avoid per-step pinned
+        # memory allocation (cudaHostAlloc). Fill the pinned CPU tensor and
+        # copy to the persistent GPU tensor with non_blocking to overlap H2D.
+        # Using a persistent tensor also avoids PT2 recompilation: since lr
+        # is the same tensor object every step, PT2 treats it as a dynamic
+        # input rather than specializing on each new tensor/value.
+        state_lists[LR_CPU_PINNED].fill_(group[LR])
+        lr = state_lists[LR_TENSOR]
+        lr.copy_(state_lists[LR_CPU_PINNED], non_blocking=True)
+        beta1 = group[BETAS][0]
+        beta3 = group[BETA3]
+        weight_decay = group[WEIGHT_DECAY]
+        peak_lr = group[PEAK_LR]
+        weight_decay_type = group[WEIGHT_DECAY_TYPE]
+        grafting_config_not_none = group[GRAFTING_CONFIG] is not None
+        perform_amortized_computation = (
+            step_val % group[PRECONDITION_FREQUENCY] == 0
+            and step_val > group[START_PRECONDITIONING_STEP]
+        ) or step_val == group[START_PRECONDITIONING_STEP]
+        use_bias_correction = group[USE_BIAS_CORRECTION]
+        # Check if we apply the grafting method or not.
+        use_grafting_method = (
+            step_val < group[START_PRECONDITIONING_STEP] and grafting_config_not_none
+        )
+        # Set train and eval interpolation coefficients if enabled.
+        # ClassicMomentumConfig is excluded from iterate_averaging_enabled: it does
+        # not use iterate averaging (no z-sequence), so weight decay shrinkage
+        # should apply to W (params).
+        iterate_averaging_config = group[ITERATE_AVERAGING_CONFIG]
+        iterate_averaging_enabled = isinstance(
+            iterate_averaging_config,
+            (GeneralizedPrimalAveragingConfig, ScheduleFreeConfig),
+        )
+        classic_momentum_enabled = isinstance(
+            iterate_averaging_config, ClassicMomentumConfig
+        )
+        if iterate_averaging_enabled:
+            train_interp_coeff, eval_interp_coeff = (
+                self._get_train_and_eval_interp_coeffs(
+                    iterate_averaging_config=iterate_averaging_config,
+                    lr=lr,
+                    state_lists=state_lists,
+                )
+            )
+        else:
+            # Coefficients unused when iterate averaging is disabled; pass the
+            # persistent 0.0 tensors as inert placeholders (PT2-stable signature).
+            train_interp_coeff = state_lists[TRAIN_INTERP_COEFF]
+            eval_interp_coeff = state_lists[EVAL_INTERP_COEFF]
+        # Extract classic momentum parameters if enabled.
+        if classic_momentum_enabled:
+            momentum_param = iterate_averaging_config.momentum
+            dampening = iterate_averaging_config.dampening
+            use_nesterov = iterate_averaging_config.use_nesterov
+        else:
+            momentum_param = 0.0
+            dampening = 0.0
+            use_nesterov = False
+
+        per_group_step(
+            state_lists,
+            step,
+            lr,
+            beta1,
+            beta3,
+            weight_decay,
+            peak_lr,
+            weight_decay_type,
+            grafting_config_not_none,
+            perform_amortized_computation,
+            use_bias_correction,
+            use_grafting_method,
+            iterate_averaging_enabled,
+            classic_momentum_enabled,
+            train_interp_coeff,
+            eval_interp_coeff,
+            momentum_param,
+            dampening,
+            use_nesterov,
+        )
+
+        # Explicitly set masked blocked gradients to None to save memory so the original param.grad has no pointer to it.
+        state_lists[MASKED_BLOCKED_GRADS] = None
+
     @overload
     @torch.no_grad()
     def step(self, closure: None = None) -> None: ...
@@ -1370,6 +1659,11 @@ def step(self, closure: Callable[[], float]) -> float: ...
     def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 
+        When ``ConcurrencyConfig.enable_cuda_stream_for_param_groups`` is set
+        and all param groups are on CUDA, each group's GPU work runs on its
+        own dedicated CUDA stream so independent NCCL communicators across
+        groups can overlap on-device.
+
         Args:
             closure (Callable[[], float] | None): A closure that reevaluates the model and returns the loss. (Default: None)
 
@@ -1384,77 +1678,30 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         for state_lists, group in zip(
             self._per_group_state_lists, self.param_groups, strict=True
         ):
-            # Construct blocked gradient list.
-            state_lists[MASKED_BLOCKED_GRADS] = state_lists[
-                DISTRIBUTOR
-            ].merge_and_block_gradients()
-
-            # Based on the current block selector, mask lists of parameters and optimizer states.
-            DistributedShampoo._mask_state_lists(
-                state_lists=state_lists,
-                group=group,
-                shampoo_pt2_enabled=self._shampoo_pt2_compile_config is not None,
-            )
-
-            # Iterate group step counter and define Python scalar step.
-            step = state_lists[STEP].add_(1)
-            step_val = step.item()
-            # NOTE: Reuse pre-allocated lr tensors to avoid per-step pinned
-            # memory allocation (cudaHostAlloc). Fill the pinned CPU tensor and
-            # copy to the persistent GPU tensor with non_blocking to overlap H2D.
-            # Using a persistent tensor also avoids PT2 recompilation: since lr
-            # is the same tensor object every step, PT2 treats it as a dynamic
-            # input rather than specializing on each new tensor/value.
-            state_lists[LR_CPU_PINNED].fill_(group[LR])
-            lr = state_lists[LR_TENSOR]
-            lr.copy_(state_lists[LR_CPU_PINNED], non_blocking=True)
-            beta1 = group[BETAS][0]
-            beta3 = group[BETA3]
-            weight_decay = group[WEIGHT_DECAY]
-            peak_lr = group[PEAK_LR]
-            weight_decay_type = group[WEIGHT_DECAY_TYPE]
-            grafting_config_not_none = group[GRAFTING_CONFIG] is not None
-            perform_amortized_computation = (
-                step_val % group[PRECONDITION_FREQUENCY] == 0
-                and step_val > group[START_PRECONDITIONING_STEP]
-            ) or step_val == group[START_PRECONDITIONING_STEP]
-            use_bias_correction = group[USE_BIAS_CORRECTION]
-            # Check if we apply the grafting method or not.
-            use_grafting_method = (
-                step_val < group[START_PRECONDITIONING_STEP]
-                and grafting_config_not_none
-            )
-            # Set train and eval interpolation coefficients if enabled.
-            train_interp_coeff, eval_interp_coeff = (
-                self._get_train_and_eval_interp_coeffs(
-                    iterate_averaging_config=group[ITERATE_AVERAGING_CONFIG],
-                    lr=lr,
-                    state_lists=state_lists,
-                )
-            )
-
-            self._per_group_step(
-                state_lists,
-                step,
-                lr,
-                beta1,
-                beta3,
-                weight_decay,
-                peak_lr,
-                weight_decay_type,
-                grafting_config_not_none,
-                perform_amortized_computation,
-                use_bias_correction,
-                use_grafting_method,
-                train_interp_coeff,
-                eval_interp_coeff,
+            stream_ctx = (
+                cuda_stream_context(state_lists[CUDA_STREAM])
+                if self._use_cuda_streams
+                else nullcontext()
             )
+            with stream_ctx:
+                self._group_step_body(state_lists, group, self._per_group_step)
 
-            # Explicitly set masked blocked gradients to None to save memory so the original param.grad has no pointer to it.
-            state_lists[MASKED_BLOCKED_GRADS] = None
+        if self._use_cuda_streams:
+            self._sync_streams_to_default()
 
         return loss
 
+    @torch.no_grad()
+    def _sync_streams_to_default(self) -> None:
+        """Make each device's default stream wait for its per-group stream so
+        subsequent default-stream work observes per-group results."""
+        for state_lists, group in zip(
+            self._per_group_state_lists, self.param_groups, strict=True
+        ):
+            torch.cuda.current_stream(group[PARAMS][0].device).wait_stream(
+                state_lists[CUDA_STREAM]
+            )
+
     # ============================================================
     # TRAIN/EVAL MODE SWITCHING
     # ============================================================
@@ -1465,12 +1712,18 @@ def train(self) -> None:
         for state_lists, group in zip(
             self._per_group_state_lists, self.param_groups, strict=True
         ):
-            # Skip groups without iterate averaging or already in train mode.
-            if group[ITERATE_AVERAGING_CONFIG] is None or state_lists[TRAIN_MODE]:
+            # Skip groups without iterate averaging, with classic momentum, or already in train mode.
+            iterate_averaging_config = group[ITERATE_AVERAGING_CONFIG]
+            if not isinstance(
+                iterate_averaging_config,
+                (GeneralizedPrimalAveragingConfig, ScheduleFreeConfig),
+            ):
+                continue
+            if state_lists[TRAIN_MODE]:
                 continue
 
             train_interp_coeff, _ = self._get_train_and_eval_interp_coeffs(
-                iterate_averaging_config=group[ITERATE_AVERAGING_CONFIG],
+                iterate_averaging_config=iterate_averaging_config,
                 lr=None,
                 state_lists=state_lists,
             )
@@ -1493,12 +1746,18 @@ def eval(self) -> None:
         for state_lists, group in zip(
             self._per_group_state_lists, self.param_groups, strict=True
         ):
-            # Skip groups without iterate averaging or already in eval mode.
-            if group[ITERATE_AVERAGING_CONFIG] is None or not state_lists[TRAIN_MODE]:
+            # Skip groups without iterate averaging, with classic momentum, or already in eval mode.
+            iterate_averaging_config = group[ITERATE_AVERAGING_CONFIG]
+            if not isinstance(
+                iterate_averaging_config,
+                (GeneralizedPrimalAveragingConfig, ScheduleFreeConfig),
+            ):
+                continue
+            if not state_lists[TRAIN_MODE]:
                 continue
 
             train_interp_coeff, _ = self._get_train_and_eval_interp_coeffs(
-                iterate_averaging_config=group[ITERATE_AVERAGING_CONFIG],
+                iterate_averaging_config=iterate_averaging_config,
                 lr=None,
                 state_lists=state_lists,
             )
@@ -1506,7 +1765,10 @@ def eval(self) -> None:
                 state_lists[WEIGHT_BUFFER_LIST],
                 state_lists[DISTRIBUTOR].local_blocked_params,
             )
-            torch._foreach_mul_(parameter_updates, 1 - 1 / train_interp_coeff)
+            torch._foreach_mul_(
+                parameter_updates,
+                1 - torch.reciprocal(train_interp_coeff),
+            )
             state_lists[DISTRIBUTOR].update_params(
                 blocked_search_directions=tuple(parameter_updates),
                 use_masked_tensors=False,
@@ -1628,8 +1890,9 @@ def _pre_load_state_dict_hook(optimizer: Optimizer, state_dict: StateDict) -> No
                 strict=True,
             )
             if group[ITERATE_AVERAGING_CONFIG] is not None
+            and not isinstance(group[ITERATE_AVERAGING_CONFIG], ClassicMomentumConfig)
         ]
-        optimizer._pre_load_train_modes = saved_train_modes  # type: ignore[attr-defined]
+        optimizer._pre_load_train_modes = saved_train_modes  # type: ignore
 
     @staticmethod
     def _post_load_state_dict_hook(optimizer: Optimizer) -> None:
diff --git a/distributed_shampoo/distributor/__init__.py b/distributed_shampoo/distributor/__init__.py
index e69de29..9071dd2 100644
--- a/distributed_shampoo/distributor/__init__.py
+++ b/distributed_shampoo/distributor/__init__.py
@@ -0,0 +1,8 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
diff --git a/distributed_shampoo/distributor/_shampoo_fully_shard_lossless_distributor.py b/distributed_shampoo/distributor/_shampoo_fully_shard_lossless_distributor.py
index b378c1c..d0e2759 100644
--- a/distributed_shampoo/distributor/_shampoo_fully_shard_lossless_distributor.py
+++ b/distributed_shampoo/distributor/_shampoo_fully_shard_lossless_distributor.py
@@ -22,9 +22,9 @@
     PARAMS,
     ShampooRuntimeConfig,
 )
-from distributed_shampoo.utils.shampoo_utils import (
-    prepare_update_param_buffers,
-    redistribute_and_update_params,
+from distributed_shampoo.utils.shampoo_fully_shard_utils import (
+    GatherGradientsContext,
+    RedistributeParamsContext,
 )
 from torch import distributed as dist, Tensor
 
@@ -67,43 +67,48 @@ def __init__(
         self._group_rank: int = dist.get_rank(group=self._dist_group)
 
         def should_assign_param_idx(i: int) -> bool:
-            if (
-                self._param_assignment_strategy
+            return (
+                i % self._group_size == self._group_rank
+                if self._param_assignment_strategy
                 == FSDPParamAssignmentStrategy.ROUND_ROBIN
-            ):
-                return i % self._group_size == self._group_rank
-            return True
+                else True
+            )
 
         self._assigned_params_mask: tuple[bool, ...] = tuple(
             should_assign_param_idx(idx) for idx in range(len(param_group[PARAMS]))
         )
 
-        # Collects and stores the model parameters assigned to this rank.
-        # Note that we explicitly disable the unnecessary gradient tracking for the all-gather collectives
-        # used to initialize the full parameters.
-        with (
-            torch.no_grad(),
-            shampoo_comm_profiler(f"{self.__class__.__name__}::full_tensor_calls"),
-        ):
-            full_params: list[Tensor] = [p.full_tensor() for p in param_group[PARAMS]]
-
-        # TODO (irisz): eagerly initialize the _assigned_full_params cannot handle dead layers parameters correctly,
-        # as we do not have gradient information during initialization. We need a way to handle dead layers parameters
-        # before doing performance optimization on the full_tensor call above (e.g. change full_tensor call to all_to_all).
-        self._assigned_full_params: list[Tensor] = [
-            p
-            for p, assigned in zip(full_params, self._assigned_params_mask)
-            if assigned
-        ]
+        # For ROUND_ROBIN strategy, create optimized redistribution context
+        # that precomputes all static information for efficient all_to_all
+        self._redistribute_ctx: RedistributeParamsContext | None = (
+            RedistributeParamsContext(
+                params=param_group[PARAMS],
+                assigned_params_mask=self._assigned_params_mask,
+                dist_group=self._dist_group,
+            )
+            if self._param_assignment_strategy
+            == FSDPParamAssignmentStrategy.ROUND_ROBIN
+            else None
+        )
 
-        # For ROUND_ROBIN strategy, creates a buffer for receiving the updated param shards.
-        self._update_param_buffers: list[Tensor] | None = (
-            prepare_update_param_buffers(param_group[PARAMS], self._group_size)
+        # For ROUND_ROBIN strategy, create gradient gathering context
+        # that uses all_to_all instead of per-param full_tensor() all-gathers
+        self._gather_grads_ctx: GatherGradientsContext | None = (
+            GatherGradientsContext(
+                params=param_group[PARAMS],
+                assigned_params_mask=self._assigned_params_mask,
+                dist_group=self._dist_group,
+            )
             if self._param_assignment_strategy
             == FSDPParamAssignmentStrategy.ROUND_ROBIN
             else None
         )
 
+        # Set _param_group early so _get_assigned_full_params can access it
+        # (super().__init__ will set it again to the same value).
+        self._param_group = param_group
+        self._assigned_full_params: list[Tensor] = self._get_assigned_full_params()
+
         super().__init__(param_group, runtime_config)
 
         if logger.isEnabledFor(logging.DEBUG):
@@ -116,6 +121,27 @@ def should_assign_param_idx(i: int) -> bool:
                 f"Local blocked params[size={len(local_block_id_list)}]: {local_block_id_list}"
             )
 
+    @torch.no_grad()
+    def _get_assigned_full_params(self) -> list[Tensor]:
+        """Gather full parameter tensors for params assigned to this rank.
+
+        For ROUND_ROBIN, uses all_to_all to gather only assigned params efficiently.
+        For REPLICATE, falls back to per-param full_tensor() all-gathers.
+
+        Returns:
+            List of full parameter tensors assigned to this rank.
+        """
+        if self._gather_grads_ctx is not None:
+            all_params = self._gather_grads_ctx.gather_params()
+            return [p for p in all_params if p is not None]
+        else:
+            full_params = [p.full_tensor() for p in self._param_group[PARAMS]]
+            return [
+                p
+                for p, assigned in zip(full_params, self._assigned_params_mask)
+                if assigned
+            ]
+
     @overload
     @torch.no_grad()
     def _get_params_or_grads(
@@ -139,13 +165,23 @@ def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None
         """
 
         if get_grad:
-            # Getting grads at every optimizer step triggers implicit all-gather. Note that p.numel()
-            # returns total number of elements in the tensor (as opposed to local shard of DTensor).
-            with shampoo_comm_profiler(f"{self.__class__.__name__}::grad_full_tensors"):
-                full_grads = [
-                    None if p.grad is None else p.grad.full_tensor()
-                    for p in self._param_group[PARAMS]
-                ]
+            gather_grads_ctx = self._gather_grads_ctx
+            if gather_grads_ctx is not None:
+                # ROUND_ROBIN: use all_to_all to gather grads efficiently.
+                # Each rank only receives full grads for its assigned params.
+                with shampoo_comm_profiler(
+                    f"{self.__class__.__name__}::gather_gradients_all_to_all"
+                ):
+                    full_grads = gather_grads_ctx.gather_gradients()
+            else:
+                # REPLICATE: fall back to per-param full_tensor() all-gathers.
+                with shampoo_comm_profiler(
+                    f"{self.__class__.__name__}::grad_full_tensors"
+                ):
+                    full_grads = [
+                        None if p.grad is None else p.grad.full_tensor()
+                        for p in self._param_group[PARAMS]
+                    ]
             return (
                 full_grad
                 for full_grad, assigned in zip(
@@ -183,13 +219,11 @@ def update_params(
         # corresponding local parameter `local_param` in the param group.
         if self._param_assignment_strategy == FSDPParamAssignmentStrategy.ROUND_ROBIN:
             with shampoo_comm_profiler(
-                f"{self.__class__.__name__}::redistribute_and_update_params"
+                f"{self.__class__.__name__}::_redistribute_ctx.redistribute_and_update_params"
             ):
-                redistribute_and_update_params(
-                    self._param_group[PARAMS],
+                assert self._redistribute_ctx is not None
+                self._redistribute_ctx.redistribute_and_update_params(
                     self._assigned_full_params,
-                    self._update_param_buffers,  # type: ignore
-                    self._dist_group,
                 )
 
         elif self._param_assignment_strategy == FSDPParamAssignmentStrategy.REPLICATE:
@@ -242,12 +276,7 @@ def refresh_assigned_full_params(self) -> None:
         1. _assigned_full_params - the cached full tensor copies
         2. _global_blocked_params and _local_blocked_params - views of the full tensors
         """
-        full_params: list[Tensor] = [p.full_tensor() for p in self._param_group[PARAMS]]
-        self._assigned_full_params = [
-            p
-            for p, assigned in zip(full_params, self._assigned_params_mask)
-            if assigned
-        ]
+        self._assigned_full_params = self._get_assigned_full_params()
         # Also re-block the parameters since _global_blocked_params are views of the old
         # _assigned_full_params tensors
         self._merge_and_block_parameters()
diff --git a/distributed_shampoo/distributor/_shampoo_hybrid_shard_lossless_distributor.py b/distributed_shampoo/distributor/_shampoo_hybrid_shard_lossless_distributor.py
index adf4467..1287408 100644
--- a/distributed_shampoo/distributor/_shampoo_hybrid_shard_lossless_distributor.py
+++ b/distributed_shampoo/distributor/_shampoo_hybrid_shard_lossless_distributor.py
@@ -7,12 +7,15 @@
 
 """
 
+import copy
+import itertools
 import logging
 from collections.abc import Iterable
 from typing import Any, Literal, overload
 
 import torch
 from distributed_shampoo.distributor.shampoo_block_info import DTensorBlockInfo
+from distributed_shampoo.distributor.shampoo_dist_utils import shampoo_comm_profiler
 from distributed_shampoo.distributor.shampoo_hybrid_shard_distributor import (
     HybridShardDistributor,
 )
@@ -23,13 +26,14 @@
     PARAMS,
     ShampooRuntimeConfig,
 )
-from distributed_shampoo.utils.shampoo_utils import (
-    compress_list,
-    prepare_update_param_buffers,
-    redistribute_and_update_params,
+from distributed_shampoo.utils.shampoo_fully_shard_utils import (
+    GatherGradientsContext,
+    RedistributeParamsContext,
 )
+from distributed_shampoo.utils.shampoo_utils import compress_list
 from torch import distributed as dist, Tensor
 from torch.distributed import tensor as dtensor
+from torch.distributed.device_mesh import init_device_mesh
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -48,6 +52,14 @@ class HybridShardLosslessDistributor(HybridShardDistributor):
 
     """
 
+    # Generates unique DeviceMesh dim names per instance. With num_sub_groups > 1,
+    # each distributor must have its own NCCL process groups to avoid deadlocks
+    # from concurrent collectives on separate CUDA streams. Unique dim names
+    # (e.g., "0_replicate", "1_replicate") ensure each init_device_mesh() call
+    # creates fresh PGs. itertools.count is atomic in CPython, so concurrent
+    # constructions get distinct indices without an explicit lock.
+    _instance_counter: "itertools.count[int]" = itertools.count()
+
     def __init__(
         self,
         param_group: dict[str, Any],
@@ -63,49 +75,78 @@ def __init__(
             f"Shampoo HybridShardLosslessDistributor {self._param_assignment_strategy=}",
         )
 
-        self._shard_group_size: int = distributed_config.device_mesh.size(1)
+        # Create a fresh DeviceMesh with unique dim names to get dedicated
+        # process groups per distributor instance. With num_sub_groups > 1,
+        # multiple distributors run concurrent collectives on separate CUDA
+        # streams — shared PGs would deadlock.
+        idx = next(HybridShardLosslessDistributor._instance_counter)
+        mesh = distributed_config.device_mesh
+        mesh_dim_names = mesh.mesh_dim_names
+        assert mesh_dim_names is not None, "DeviceMesh must have mesh_dim_names"
+        self._device_mesh: torch.distributed.device_mesh.DeviceMesh = init_device_mesh(
+            mesh.device_type,
+            tuple(mesh.mesh.shape),
+            mesh_dim_names=tuple(f"{idx}_{name}" for name in mesh_dim_names),
+        )
+        # Shallow-copy the config to give the parent the fresh device mesh
+        # without mutating the caller's original config.
+        distributed_config = copy.copy(distributed_config)
+        distributed_config.device_mesh = self._device_mesh
+        param_group[DISTRIBUTED_CONFIG] = distributed_config
+
+        self._shard_group_size: int = self._device_mesh.size(1)
         # Initialize distributed group for communicating across the shard dimension.
-        self._shard_dist_group: dist.ProcessGroup = (
-            distributed_config.device_mesh.get_group(mesh_dim=1)
+        self._shard_dist_group: dist.ProcessGroup = self._device_mesh.get_group(
+            mesh_dim=1
         )
         self._shard_group_rank: int = dist.get_rank(self._shard_dist_group)
 
         # Stores full parameters (as opposed to DTensors) for the model parameters assigned to this rank.
         # For example, when the strategy is REPLICATE, it stores the full parameters on all ranks.
         def should_assign_param_idx(i: int) -> bool:
-            if (
-                self._param_assignment_strategy
+            return (
+                i % self._shard_group_size == self._shard_group_rank
+                if self._param_assignment_strategy
                 == FSDPParamAssignmentStrategy.ROUND_ROBIN
-            ):
-                return i % self._shard_group_size == self._shard_group_rank
-            return True
+                else True
+            )
 
         with torch.no_grad():
             self._assigned_params_mask: tuple[bool, ...] = tuple(
                 should_assign_param_idx(idx) for idx in range(len(param_group[PARAMS]))
             )
 
-        # Collects and stores the model parameters assigned to this rank.
-        with torch.no_grad():
-            full_params: list[Tensor] = [p.full_tensor() for p in param_group[PARAMS]]
-
-        # TODO (irisz): eagerly initialize the _assigned_full_params cannot handle dead layers parameters correctly,
-        # as we do not have gradient information during initialization. We need a way to handle dead layers parameters
-        # before doing performance optimization on the full_tensor call above (e.g. change full_tensor call to all_to_all).
-        self._assigned_full_params: list[Tensor] = [
-            p
-            for p, assigned in zip(full_params, self._assigned_params_mask)
-            if assigned
-        ]
+        # For ROUND_ROBIN strategy, create optimized redistribution context
+        # that precomputes all static information for efficient all_to_all
+        self._redistribute_ctx: RedistributeParamsContext | None = (
+            RedistributeParamsContext(
+                params=param_group[PARAMS],
+                assigned_params_mask=self._assigned_params_mask,
+                dist_group=self._shard_dist_group,
+            )
+            if self._param_assignment_strategy
+            == FSDPParamAssignmentStrategy.ROUND_ROBIN
+            else None
+        )
 
-        # For ROUND_ROBIN strategy, creates a buffer for receiving the updated param shards.
-        self._update_param_buffers: list[Tensor] | None = (
-            prepare_update_param_buffers(param_group[PARAMS], self._shard_group_size)
+        # For ROUND_ROBIN strategy, create gradient gathering context
+        # that uses all_to_all instead of per-param full_tensor() all-gathers
+        self._gather_grads_ctx: GatherGradientsContext | None = (
+            GatherGradientsContext(
+                params=param_group[PARAMS],
+                assigned_params_mask=self._assigned_params_mask,
+                dist_group=self._shard_dist_group,
+            )
             if self._param_assignment_strategy
             == FSDPParamAssignmentStrategy.ROUND_ROBIN
             else None
         )
 
+        # Set _param_group early so _get_assigned_full_params can access it
+        # (super().__init__ will set it again to the same value).
+        self._param_group = param_group
+        self._assigned_full_params: list[Tensor] = self._get_assigned_full_params()
+
         super().__init__(param_group, runtime_config)
 
         if logger.isEnabledFor(logging.DEBUG):
@@ -118,6 +159,27 @@ def should_assign_param_idx(i: int) -> bool:
                 f"Local blocked params[size={len(local_block_id_list)}]: {local_block_id_list}"
             )
 
+    @torch.no_grad()
+    def _get_assigned_full_params(self) -> list[Tensor]:
+        """Gather full parameter tensors for params assigned to this rank.
+
+        For ROUND_ROBIN, uses all_to_all to gather only assigned params efficiently.
+        For REPLICATE, falls back to per-param full_tensor() all-gathers.
+
+        Returns:
+            List of full parameter tensors assigned to this rank.
+        """
+        if self._gather_grads_ctx is not None:
+            all_params = self._gather_grads_ctx.gather_params()
+            return [p for p in all_params if p is not None]
+        else:
+            full_params = [p.full_tensor() for p in self._param_group[PARAMS]]
+            return [
+                p
+                for p, assigned in zip(full_params, self._assigned_params_mask)
+                if assigned
+            ]
+
     def _get_composable_block_id_rank(self) -> int | None:
         """For lossless shampoo distributor, it's unnecessary to include rank in block id."""
         return None
@@ -144,11 +206,22 @@ def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None
             local (Iterable[Tensor | None]): assigned params (or grad) from the param_group.
         """
         if get_grad:
-            # NOTE: getting grads at every optimizer step triggers implicit all-gather.
-            full_grads = (
-                None if p.grad is None else p.grad.full_tensor()
-                for p in self._param_group[PARAMS]
-            )
+            gather_grads_ctx = self._gather_grads_ctx
+            if gather_grads_ctx is not None:
+                # ROUND_ROBIN: use all_to_all to gather grads efficiently.
+                with shampoo_comm_profiler(
+                    f"{self.__class__.__name__}::gather_gradients_all_to_all"
+                ):
+                    full_grads = gather_grads_ctx.gather_gradients()
+            else:
+                # REPLICATE: fall back to per-param full_tensor() all-gathers.
+                with shampoo_comm_profiler(
+                    f"{self.__class__.__name__}::grad_full_tensors"
+                ):
+                    full_grads = [
+                        None if p.grad is None else p.grad.full_tensor()
+                        for p in self._param_group[PARAMS]
+                    ]
             return (
                 full_grad
                 for full_grad, assigned in zip(
@@ -181,12 +254,13 @@ def update_params(
         # redistribute it according to the device mesh to get the locally assigned slice, and copy the slice to the
         # corresponding local parameter `local_param` in the param group.
         if self._param_assignment_strategy == FSDPParamAssignmentStrategy.ROUND_ROBIN:
-            redistribute_and_update_params(
-                self._param_group[PARAMS],
-                self._assigned_full_params,
-                self._update_param_buffers,  # type: ignore
-                self._shard_dist_group,
-            )
+            with shampoo_comm_profiler(
+                f"{self.__class__.__name__}::_redistribute_ctx.redistribute_and_update_params"
+            ):
+                assert self._redistribute_ctx is not None
+                self._redistribute_ctx.redistribute_and_update_params(
+                    self._assigned_full_params,
+                )
 
         # Copy the updated full parameters to the original parameters.
         elif self._param_assignment_strategy == FSDPParamAssignmentStrategy.REPLICATE:
@@ -244,12 +318,7 @@ def refresh_assigned_full_params(self) -> None:
         2. _global_blocked_params and _local_blocked_params - views of the full tensors
         3. _global_masked_blocked_params - HybridShardDistributor's reference to global blocked params
         """
-        full_params: list[Tensor] = [p.full_tensor() for p in self._param_group[PARAMS]]
-        self._assigned_full_params = [
-            p
-            for p, assigned in zip(full_params, self._assigned_params_mask)
-            if assigned
-        ]
+        self._assigned_full_params = self._get_assigned_full_params()
         # Also re-block the parameters since _global_blocked_params are views of the old
         # _assigned_full_params tensors
         self._merge_and_block_parameters()
diff --git a/distributed_shampoo/distributor/gpu_tests/distributor_test_utils.py b/distributed_shampoo/distributor/gpu_tests/distributor_test_utils.py
index 256419e..268c632 100644
--- a/distributed_shampoo/distributor/gpu_tests/distributor_test_utils.py
+++ b/distributed_shampoo/distributor/gpu_tests/distributor_test_utils.py
@@ -9,14 +9,102 @@
 
 import abc
 import unittest
+from collections.abc import Callable
+from functools import partial
 
 import torch
+from distributed_shampoo.distributed_shampoo import DistributedShampoo
 from distributed_shampoo.distributor.shampoo_block_info import (
     BlockInfo,
     DTensorBlockInfo,
 )
 from distributed_shampoo.distributor.shampoo_distributor import DistributorInterface
+from distributed_shampoo.shampoo_types import (
+    AdaGradPreconditionerConfig,
+    DistributedConfig,
+    GeneralizedPrimalAveragingConfig,
+    ShampooPT2CompileConfig,
+    WeightDecayType,
+)
 from torch import nn
+from torch.optim.optimizer import ParamsT
+
+# Shared constants for fully-shard and hybrid-shard param group tests.
+PRECONDITIONER_DIM = 4
+
+FULLY_SHARD_TEST_MODEL_LAYER_DIMS: tuple[tuple[int, ...], ...] = (
+    (4 * PRECONDITIONER_DIM, 2 * PRECONDITIONER_DIM, 1),
+    (3 * PRECONDITIONER_DIM - 1, PRECONDITIONER_DIM + 1, PRECONDITIONER_DIM - 1),
+    (2, 2 * PRECONDITIONER_DIM - 1),
+    (PRECONDITIONER_DIM, 1),
+)
+
+FULLY_SHARD_DEAD_MODEL_LAYER_DIMS: tuple[tuple[int, ...], ...] = (
+    (PRECONDITIONER_DIM, 1),
+    (),
+)
+
+HYBRID_SHARD_TEST_MODEL_LAYER_DIMS: tuple[tuple[int, ...], ...] = (
+    (4 * PRECONDITIONER_DIM, 2 * PRECONDITIONER_DIM, 1, 1, 1),
+    (3 * PRECONDITIONER_DIM - 1, PRECONDITIONER_DIM + 1, PRECONDITIONER_DIM - 1, 1, 1),
+    (PRECONDITIONER_DIM, 1),
+)
+
+HYBRID_SHARD_DEAD_MODEL_LAYER_DIMS: tuple[tuple[int, ...], ...] = (
+    (PRECONDITIONER_DIM, 1),
+)
+
+# Model dims for param group tests. Uses a single fixed model with enough
+# params for num_sub_groups=4. No dims parametrization needed — the model
+# size doesn't affect the splitting/threading logic.
+# FSDP: 5 params (4 live + 1 dead). No shard_size constraint.
+PARAM_GROUP_TEST_MODEL_DIMS: tuple[int, ...] = (
+    4 * PRECONDITIONER_DIM,
+    2 * PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    1,
+)
+PARAM_GROUP_TEST_DEAD_DIMS: tuple[int, ...] = (PRECONDITIONER_DIM, 1)
+# HSDP with (2,2) mesh: shard_size=2, so each sub-group needs >= 2 params.
+# num_sub_groups=4 needs >= 8 params. This model has 8 live + 1 dead = 9.
+# Two wider blocks (2*PD) mixed with uniform PD layers give three numel tiers
+# (32 / 16 / 4 with PD=4) so greedy LPT bin-packing has to interleave rather
+# than trivially round-robin; no single weight dominates enough to strand a bin
+# and trip the per-bin shard_size guard in _create_sub_groups.
+HYBRID_SHARD_PARAM_GROUP_TEST_MODEL_DIMS: tuple[int, ...] = (
+    2 * PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    2 * PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    PRECONDITIONER_DIM,
+    1,
+)
+
+
+def shampoo_optim_factory(
+    distributed_config: DistributedConfig,
+    shampoo_pt2_compile_config: ShampooPT2CompileConfig | None = None,
+    start_preconditioning_step: int = 2,
+) -> Callable[[ParamsT], torch.optim.Optimizer]:
+    """Shared optimizer factory for FSDP/HSDP lossless distributor tests."""
+    return partial(
+        DistributedShampoo,
+        lr=0.001,
+        betas=(0.9, 1.0),
+        epsilon=1e-8,
+        weight_decay=0.0,
+        max_preconditioner_dim=PRECONDITIONER_DIM,
+        precondition_frequency=1,
+        start_preconditioning_step=start_preconditioning_step,
+        weight_decay_type=WeightDecayType.DECOUPLED,
+        grafting_config=AdaGradPreconditionerConfig(epsilon=1e-8),
+        distributed_config=distributed_config,
+        iterate_averaging_config=GeneralizedPrimalAveragingConfig(),
+        shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+    )
 
 
 class DistributorOnEmptyParamTest:
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py
index 5c469b8..ab34994 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_checkpoint_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import copy
 from collections.abc import Callable
 from functools import partial
@@ -39,7 +37,7 @@
     StateDictOptions,
 )
 from torch.distributed.device_mesh import init_device_mesh
-from torch.distributed.fsdp import fully_shard
+from torch.distributed.fsdp import FSDPModule, fully_shard
 from torch.distributed.tensor import DTensor
 from torch.optim.optimizer import ParamsT
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
@@ -261,8 +259,10 @@ def world_size(self) -> int:
     def _construct_model(
         model_linear_layers_dims: tuple[int, ...],
         model_dead_layers_dims: tuple[int, ...],
-        post_model_decoration: Callable[[nn.Module], nn.Module] = lambda x: x,
-    ) -> nn.Module:
+        post_model_decoration: Callable[
+            [nn.Module], nn.Module | FSDPModule
+        ] = lambda x: x,
+    ) -> nn.Module | FSDPModule:
         # Using partial here to prevent Pyre complain on incompatible parameter type.
         model, _, _, _ = partial(
             construct_training_problem, post_model_decoration=post_model_decoration
@@ -373,8 +373,12 @@ def test_checkpoint_preemption_bitwise_equivalence(
         model1 = self._construct_model(
             model_linear_layers_dims=model_linear_layers_dims,
             model_dead_layers_dims=model_dead_layers_dims,
-            post_model_decoration=post_model_decoration,  # type: ignore[arg-type]
+            post_model_decoration=post_model_decoration,
         )
+        # FSDPModule is a runtime mixin on nn.Module, so this assert always
+        # holds; it narrows the Union for the downstream PyTorch APIs that
+        # expect plain nn.Module.
+        assert isinstance(model1, nn.Module)
         input_dim = model_linear_layers_dims[0]
         optimizer1 = self._shampoo_optim_factory(distributed_config=config)(
             model1.parameters()
@@ -408,8 +412,9 @@ def test_checkpoint_preemption_bitwise_equivalence(
         model2 = self._construct_model(
             model_linear_layers_dims=model_linear_layers_dims,
             model_dead_layers_dims=model_dead_layers_dims,
-            post_model_decoration=post_model_decoration,  # type: ignore[arg-type]
+            post_model_decoration=post_model_decoration,
         )
+        assert isinstance(model2, nn.Module)
         # Use tiny lr for init step to minimize impact
         optimizer2 = self._shampoo_optim_factory(distributed_config=config, lr=1e-10)(
             model2.parameters()
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py
index 676a5ce..1c8bd0b 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_ddp_distributor_test.py
@@ -7,12 +7,11 @@
 
 """
 
-#!/usr/bin/env python3
-
 import abc
 import os
 import unittest
 from collections.abc import Callable
+from dataclasses import replace
 from functools import partial
 from typing import Any, cast, Dict
 
@@ -25,6 +24,7 @@
 from distributed_shampoo.distributor.shampoo_ddp_distributor import DDPDistributor
 from distributed_shampoo.shampoo_types import (
     AdaGradPreconditionerConfig,
+    ClassicMomentumConfig,
     DDPDistributedConfig,
     DefaultEigenvalueCorrectedShampooConfig,
     DefaultShampooConfig,
@@ -41,13 +41,16 @@
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
+from distributed_shampoo.utils.shampoo_utils import pack_upper_triangular
 from torch import distributed as dist, nn, tensor
 from torch.distributed.checkpoint.state_dict import get_optimizer_state_dict
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.placement_types import Replicate
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.optimizer import ParamsT
 from torch.testing._comparison import default_tolerances
 from torch.testing._internal.common_distributed import (
@@ -62,6 +65,25 @@
 PRECONDITIONER_DIM = 3
 
 
+def _shampoo_factor_dtensor(
+    local_tensor: torch.Tensor,
+    device_mesh: DeviceMesh,
+    use_symmetric_packing: bool,
+) -> DTensor:
+    """Builds an expected Shampoo factor / inverse factor matrix DTensor.
+
+    Packs `local_tensor` to upper triangular format when symmetric packing is enabled.
+    Shampoo state in this test always uses Replicate placement.
+    """
+    return DTensor.from_local(
+        local_tensor=pack_upper_triangular(local_tensor)
+        if use_symmetric_packing
+        else local_tensor,
+        device_mesh=device_mesh,
+        placements=(Replicate(),),
+    )
+
+
 # Use outer class as wrapper to avoid running the abstract test.
 class AbstractTest:
     @instantiate_parametrized_tests
@@ -93,17 +115,19 @@ def _shampoo_optim_factory(
             distributed_config: DDPDistributedConfig | SingleDeviceDistributedConfig,
             preconditioner_config: PreconditionerConfig = DefaultShampooConfig,
             iterate_averaging_config: IterateAveragingConfig | None = None,
+            weight_decay: float = 0.0,
+            weight_decay_type: WeightDecayType = WeightDecayType.DECOUPLED,
         ) -> Callable[[ParamsT], torch.optim.Optimizer]:
             return partial(
                 DistributedShampoo,
                 lr=0.001,
                 betas=(0.9, 1.0),
                 epsilon=1e-8,
-                weight_decay=0.0,
+                weight_decay=weight_decay,
                 max_preconditioner_dim=PRECONDITIONER_DIM,
                 precondition_frequency=1,
                 start_preconditioning_step=2,
-                weight_decay_type=WeightDecayType.DECOUPLED,
+                weight_decay_type=weight_decay_type,
                 grafting_config=AdaGradPreconditionerConfig(epsilon=1e-8),
                 distributed_config=distributed_config,
                 preconditioner_config=preconditioner_config,
@@ -186,6 +210,79 @@ def test_losses(
                 fill=0.01,
                 rtol=rtol,
                 atol=atol,
+                rank=self.rank,
+                world_size=self.world_size,
+                # DDP wrapping is needed so gradients are all-reduced across ranks,
+                # matching real DDP training semantics. find_unused_parameters=True
+                # is required because the model's dead layers don't participate in
+                # forward() and thus never receive gradients.
+                experimental_post_model_decoration=partial(
+                    DDP,
+                    device_ids=[self.rank] if self._device.type == "cuda" else None,
+                    find_unused_parameters=True,
+                ),
+            )
+
+        @parametrize(
+            "weight_decay_type",
+            (
+                WeightDecayType.DECOUPLED,
+                WeightDecayType.CORRECTED,
+                WeightDecayType.INDEPENDENT,
+                WeightDecayType.L2,
+            ),
+        )
+        @parametrize("num_trainers_per_group", (-1, 1, 2))
+        @parametrize(
+            "iterate_averaging_config",
+            (
+                None,
+                GeneralizedPrimalAveragingConfig(),
+                ScheduleFreeConfig(),
+                ClassicMomentumConfig(),
+            ),
+        )
+        def test_weight_decay_losses(
+            self,
+            num_trainers_per_group: int,
+            weight_decay_type: WeightDecayType,
+            iterate_averaging_config: IterateAveragingConfig | None,
+        ) -> None:
+            self._init_distributed()
+
+            compare_two_optimizers_on_weight_and_loss(
+                control_optim_factory=self._shampoo_optim_factory(
+                    distributed_config=DefaultSingleDeviceDistributedConfig,
+                    iterate_averaging_config=iterate_averaging_config,
+                    weight_decay=0.3,
+                    weight_decay_type=weight_decay_type,
+                ),
+                experimental_optim_factory=self._shampoo_optim_factory(
+                    distributed_config=DDPDistributedConfig(
+                        communication_dtype=torch.float32,
+                        num_trainers_per_group=num_trainers_per_group,
+                    ),
+                    iterate_averaging_config=iterate_averaging_config,
+                    weight_decay=0.3,
+                    weight_decay_type=weight_decay_type,
+                ),
+                model_linear_layers_dims=(
+                    PRECONDITIONER_DIM * 4,
+                    PRECONDITIONER_DIM * 2,
+                    1,
+                ),
+                model_dead_layers_dims=(PRECONDITIONER_DIM, PRECONDITIONER_DIM),
+                device=self._device,
+                fill=0.01,
+                rtol=0.0,
+                atol=0.0,
+                rank=self.rank,
+                world_size=self.world_size,
+                experimental_post_model_decoration=partial(
+                    DDP,
+                    device_ids=[self.rank] if self._device.type == "cuda" else None,
+                    find_unused_parameters=True,
+                ),
             )
 
         @parametrize(
@@ -254,15 +351,34 @@ def test_can_run_with_empty_local_params(
                 fill=0.01,
                 rtol=rtol,
                 atol=atol,
+                rank=self.rank,
+                world_size=self.world_size,
+                # DDP wrapping with find_unused_parameters=True (see test_losses comment).
+                experimental_post_model_decoration=partial(
+                    DDP,
+                    device_ids=[self.rank] if self._device.type == "cuda" else None,
+                    find_unused_parameters=True,
+                ),
             )
 
-        def test_state_dict(self) -> None:
+        @parametrize("use_symmetric_packing", (False, True))
+        def test_state_dict(self, use_symmetric_packing: bool) -> None:
             self._init_distributed()
 
             num_steps = 3
+            global_train_data = generate_global_train_data(
+                num_steps=num_steps,
+                world_size=self.world_size,
+                data_shape=(PRECONDITIONER_DIM * 2,),
+                device=self._device,
+            )
             model, _, _, _, optimizer = train_model(
                 optim_factory=AbstractTest.ShampooDDPDistributorDeviceTest._shampoo_optim_factory(
-                    distributed_config=DDPDistributedConfig()
+                    distributed_config=DDPDistributedConfig(),
+                    preconditioner_config=replace(
+                        DefaultShampooConfig,
+                        use_symmetric_packing=use_symmetric_packing,
+                    ),
                 ),
                 # Setting model_linear_layers_dims to creates an model with one linear layer with (PRECONDITIONER_DIM * 2)xPRECONDITIONER_DIM weight.
                 # Because Shampoo's max_preconditioner_dim = PRECONDITIONER_DIM, there will be two blocks; rank 0 has block 0 and rank 1 has block 1.
@@ -277,6 +393,7 @@ def test_state_dict(self) -> None:
                     fill=0.01,
                 ),
                 num_steps=num_steps,
+                train_data=global_train_data[:, self.rank],
             )
 
             assert isinstance(optimizer, DistributedShampoo)
@@ -293,6 +410,9 @@ def test_state_dict(self) -> None:
                     1,
                 ],
             )
+            factor = partial(
+                _shampoo_factor_dtensor, use_symmetric_packing=use_symmetric_packing
+            )
 
             # Define the expected state dictionary for each rank.
             # The state_dict is keyed by parameter index (0 and 1).
@@ -303,27 +423,25 @@ def test_state_dict(self) -> None:
                         "block_0": {
                             "shampoo": {
                                 "factor_matrices": {
-                                    0: DTensor.from_local(
-                                        local_tensor=tensor([[0.0016058803303167224]]),
+                                    0: factor(
+                                        local_tensor=tensor([[0.0001959779328899458]]),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                                 "inv_factor_matrices": {
-                                    0: DTensor.from_local(
-                                        local_tensor=tensor([[24.9541072845459]]),
+                                    0: factor(
+                                        local_tensor=tensor([[71.43077087402344]]),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                             },
                             "adagrad": DTensor.from_local(
-                                local_tensor=tensor([0.0016058803303167224]),
+                                local_tensor=tensor([0.0001959779328899458]),
                                 device_mesh=mesh_0,
                                 placements=(Replicate(),),
                             ),
                             "filtered_grad": DTensor.from_local(
-                                local_tensor=tensor([0.0037594244349747896]),
+                                local_tensor=tensor([0.0018590810941532254]),
                                 device_mesh=mesh_0,
                                 placements=(Replicate(),),
                             ),
@@ -334,99 +452,95 @@ def test_state_dict(self) -> None:
                         "block_1": {
                             "shampoo": {
                                 "factor_matrices": {
-                                    0: DTensor.from_local(
+                                    0: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
+                                                    1.361092654406093e-05,
+                                                    1.361092654406093e-05,
+                                                    1.361092654406093e-05,
                                                 ],
                                                 [
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
+                                                    1.361092654406093e-05,
+                                                    1.3610928363050334e-05,
+                                                    1.361092654406093e-05,
                                                 ],
                                                 [
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
-                                                    6.645893154200166e-05,
+                                                    1.361092654406093e-05,
+                                                    1.361092654406093e-05,
+                                                    1.361092654406093e-05,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
-                                    1: DTensor.from_local(
+                                    1: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    4.689853813033551e-05,
-                                                    -3.546147490851581e-05,
-                                                    -3.7919791793683544e-05,
+                                                    1.2197860996820964e-05,
+                                                    1.3136921097611776e-06,
+                                                    -1.6640237845422234e-06,
                                                 ],
                                                 [
-                                                    -3.546147490851581e-05,
-                                                    3.5329150705365464e-05,
-                                                    1.7933603885467164e-05,
+                                                    1.3136921097611776e-06,
+                                                    7.509043371101143e-06,
+                                                    -1.2525374586402904e-05,
                                                 ],
                                                 [
-                                                    -3.7919791793683544e-05,
-                                                    1.7933603885467164e-05,
-                                                    0.0001171490948763676,
+                                                    -1.6640237845422234e-06,
+                                                    -1.2525374586402904e-05,
+                                                    2.112587753799744e-05,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                                 "inv_factor_matrices": {
-                                    0: DTensor.from_local(
+                                    0: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    69.4718017578125,
-                                                    -30.528188705444336,
-                                                    -30.52819061279297,
+                                                    70.83631896972656,
+                                                    -29.16368293762207,
+                                                    -29.16368293762207,
                                                 ],
                                                 [
-                                                    -30.528188705444336,
-                                                    69.45362091064453,
-                                                    -30.510019302368164,
+                                                    -29.163679122924805,
+                                                    70.83528900146484,
+                                                    -29.162641525268555,
                                                 ],
                                                 [
-                                                    -30.52819061279297,
-                                                    -30.510019302368164,
-                                                    69.45362854003906,
+                                                    -29.163679122924805,
+                                                    -29.16264533996582,
+                                                    70.83528900146484,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
-                                    1: DTensor.from_local(
+                                    1: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    16.394943237304688,
-                                                    5.6407470703125,
-                                                    1.8611559867858887,
+                                                    16.980432510375977,
+                                                    -1.1272867918014526,
+                                                    -0.1816159188747406,
                                                 ],
                                                 [
-                                                    5.6407470703125,
-                                                    16.996492385864258,
-                                                    0.27824175357818604,
+                                                    -1.1272867918014526,
+                                                    49.9351692199707,
+                                                    21.47319793701172,
                                                 ],
                                                 [
-                                                    1.8611557483673096,
-                                                    0.27824151515960693,
-                                                    10.019035339355469,
+                                                    -0.1816159337759018,
+                                                    21.47319793701172,
+                                                    26.421972274780273,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_0,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                             },
@@ -434,19 +548,19 @@ def test_state_dict(self) -> None:
                                 local_tensor=tensor(
                                     [
                                         [
-                                            1.563284604344517e-05,
-                                            1.1776383871620055e-05,
-                                            3.904970071744174e-05,
+                                            4.065953362442087e-06,
+                                            2.503014229660039e-06,
+                                            7.04195917933248e-06,
                                         ],
                                         [
-                                            1.563284604344517e-05,
-                                            1.1776383871620055e-05,
-                                            3.904970071744174e-05,
+                                            4.065953362442087e-06,
+                                            2.5030146844073897e-06,
+                                            7.04195917933248e-06,
                                         ],
                                         [
-                                            1.563284604344517e-05,
-                                            1.1776383871620055e-05,
-                                            3.904970071744174e-05,
+                                            4.065953362442087e-06,
+                                            2.503014229660039e-06,
+                                            7.04195917933248e-06,
                                         ],
                                     ]
                                 ),
@@ -457,19 +571,19 @@ def test_state_dict(self) -> None:
                                 local_tensor=tensor(
                                     [
                                         [
-                                            -0.0003921023744624108,
-                                            0.00012148835230618715,
-                                            0.0008510660263709724,
+                                            -0.0002331818686798215,
+                                            -8.544711454305798e-05,
+                                            0.00015910383081063628,
                                         ],
                                         [
-                                            -0.0003921023744624108,
-                                            0.00012148835230618715,
-                                            0.0008510660263709724,
+                                            -0.0002331818686798215,
+                                            -8.54471290949732e-05,
+                                            0.00015910383081063628,
                                         ],
                                         [
-                                            -0.0003921023744624108,
-                                            0.00012148835230618715,
-                                            0.0008510660263709724,
+                                            -0.0002331818686798215,
+                                            -8.544711454305798e-05,
+                                            0.00015910383081063628,
                                         ],
                                     ]
                                 ),
@@ -484,99 +598,95 @@ def test_state_dict(self) -> None:
                         "block_0": {
                             "shampoo": {
                                 "factor_matrices": {
-                                    0: DTensor.from_local(
+                                    0: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
                                                 ],
                                                 [
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
                                                 ],
                                                 [
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
-                                                    0.00011197220010217279,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
+                                                    2.4930672225309536e-05,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_1,
-                                        placements=(Replicate(),),
                                     ),
-                                    1: DTensor.from_local(
+                                    1: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    0.000132521876366809,
-                                                    0.00013956104521639645,
-                                                    -1.1767312571464572e-05,
+                                                    1.4635477782576345e-05,
+                                                    -1.1405569239286706e-05,
+                                                    -8.130889909807593e-06,
                                                 ],
                                                 [
-                                                    0.00013956104521639645,
-                                                    0.0001891223801067099,
-                                                    6.5871563492692076e-06,
+                                                    -1.1405569239286706e-05,
+                                                    4.607178925652988e-05,
+                                                    2.4604041755083017e-05,
                                                 ],
                                                 [
-                                                    -1.1767312571464572e-05,
-                                                    6.5871563492692076e-06,
-                                                    1.4272330190578941e-05,
+                                                    -8.130889909807593e-06,
+                                                    2.4604041755083017e-05,
+                                                    1.4084745998843573e-05,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_1,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                                 "inv_factor_matrices": {
-                                    0: DTensor.from_local(
+                                    0: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    69.12242889404297,
-                                                    -30.855859756469727,
-                                                    -30.88007926940918,
+                                                    70.21800231933594,
+                                                    -29.731094360351562,
+                                                    -29.73412322998047,
                                                 ],
                                                 [
-                                                    -30.855863571166992,
-                                                    69.09217071533203,
-                                                    -30.84980583190918,
+                                                    -29.731094360351562,
+                                                    70.24112701416016,
+                                                    -29.75722885131836,
                                                 ],
                                                 [
-                                                    -30.88007926940918,
-                                                    -30.849802017211914,
-                                                    69.11640167236328,
+                                                    -29.734119415283203,
+                                                    -29.75722885131836,
+                                                    70.24415588378906,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_1,
-                                        placements=(Replicate(),),
                                     ),
-                                    1: DTensor.from_local(
+                                    1: factor(
                                         local_tensor=tensor(
                                             [
                                                 [
-                                                    13.893352508544922,
-                                                    -5.123697280883789,
-                                                    4.147429466247559,
+                                                    17.4156436920166,
+                                                    0.022538283839821815,
+                                                    3.622297525405884,
                                                 ],
                                                 [
-                                                    -5.123697280883789,
-                                                    11.735774993896484,
-                                                    -3.2658231258392334,
+                                                    0.02253795601427555,
+                                                    16.994813919067383,
+                                                    -10.457884788513184,
                                                 ],
                                                 [
-                                                    4.147429466247559,
-                                                    -3.2658231258392334,
-                                                    19.706111907958984,
+                                                    3.622298002243042,
+                                                    -10.4578857421875,
+                                                    32.26253890991211,
                                                 ],
                                             ]
                                         ),
                                         device_mesh=mesh_1,
-                                        placements=(Replicate(),),
                                     ),
                                 },
                             },
@@ -584,19 +694,19 @@ def test_state_dict(self) -> None:
                                 local_tensor=tensor(
                                     [
                                         [
-                                            4.417396485223435e-05,
-                                            6.304078851826489e-05,
-                                            4.757443548442097e-06,
+                                            4.878492745774565e-06,
+                                            1.5357263691839762e-05,
+                                            4.694914878200507e-06,
                                         ],
                                         [
-                                            4.417396485223435e-05,
-                                            6.304078851826489e-05,
-                                            4.7574430936947465e-06,
+                                            4.878493200521916e-06,
+                                            1.5357263691839762e-05,
+                                            4.694915332947858e-06,
                                         ],
                                         [
-                                            4.417396485223435e-05,
-                                            6.304078851826489e-05,
-                                            4.757443548442097e-06,
+                                            4.878492745774565e-06,
+                                            1.5357263691839762e-05,
+                                            4.694914878200507e-06,
                                         ],
                                     ]
                                 ),
@@ -607,19 +717,19 @@ def test_state_dict(self) -> None:
                                 local_tensor=tensor(
                                     [
                                         [
-                                            0.000822467845864594,
-                                            0.0009405062301084399,
-                                            7.864914368838072e-05,
+                                            0.00016447670350316912,
+                                            8.83667089510709e-05,
+                                            -4.004316360806115e-05,
                                         ],
                                         [
-                                            0.000822467845864594,
-                                            0.0009405062301084399,
-                                            7.86491364124231e-05,
+                                            0.00016447667439933866,
+                                            8.836669439915568e-05,
+                                            -4.004319998784922e-05,
                                         ],
                                         [
-                                            0.000822467845864594,
-                                            0.0009405062301084399,
-                                            7.864914368838072e-05,
+                                            0.0001644766889512539,
+                                            8.836670167511329e-05,
+                                            -4.004317815997638e-05,
                                         ],
                                     ]
                                 ),
@@ -674,6 +784,12 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
             self._init_distributed()
 
             steps_without_gradients = 2
+            global_train_data = generate_global_train_data(
+                num_steps=steps_without_gradients,
+                world_size=self.world_size,
+                data_shape=(PRECONDITIONER_DIM * 4,),
+                device=self._device,
+            )
             with unittest.mock.patch.object(torch.Tensor, "backward") as mock_backward:
                 # By mocking the backward() method, we're intercepting gradient calculation.
                 # This effectively simulates running forward passes without computing gradients.
@@ -694,6 +810,7 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
                         device=self._device,
                     ),
                     num_steps=steps_without_gradients,
+                    train_data=global_train_data[:, self.rank],
                 )
 
             # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -706,6 +823,12 @@ def test_some_ranks_with_no_grads_due_to_dead_layers(
             self._init_distributed()
 
             num_steps = 3
+            global_train_data = generate_global_train_data(
+                num_steps=num_steps,
+                world_size=self.world_size,
+                data_shape=(PRECONDITIONER_DIM,),
+                device=self._device,
+            )
             model, _, _, _, optimizer = train_model(
                 optim_factory=AbstractTest.ShampooDDPDistributorDeviceTest._shampoo_optim_factory(
                     distributed_config=DDPDistributedConfig(
@@ -721,6 +844,7 @@ def test_some_ranks_with_no_grads_due_to_dead_layers(
                     device=self._device,
                 ),
                 num_steps=num_steps,
+                train_data=global_train_data[:, self.rank],
             )
 
             assert isinstance(optimizer, DistributedShampoo)
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_dist_utils_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_dist_utils_test.py
index 9adeb8b..c5c1b9d 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_dist_utils_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_dist_utils_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 from functools import partial
 from operator import attrgetter
 from unittest import mock
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_distributor_test.py
index 7d91201..4dc8cd7 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Callable
 from functools import partial
@@ -37,6 +35,7 @@
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
 from torch import distributed as dist, nn
@@ -134,6 +133,19 @@ def test_all_ranks_with_no_grads(self) -> None:
         fsdp_config = FSDPDistributedConfig(param_to_metadata={})
 
         steps_without_gradients = 2
+        model_linear_layers_dims = (
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            2 * PRECONDITIONER_DIM,
+            9,
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            1,
+        )
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=dist.get_world_size(),
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch.object(torch.Tensor, "backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
@@ -147,6 +159,7 @@ def test_all_ranks_with_no_grads(self) -> None:
                     distributed_config=fsdp_config,
                 ),
                 num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -155,6 +168,20 @@ def test_all_ranks_with_no_grads(self) -> None:
     @skip_if_lt_x_gpu(2)
     def test_fsdp_shampoo_against_default_shampoo(self) -> None:
         fsdp_config = FSDPDistributedConfig(param_to_metadata={})
+        model_linear_layers_dims = (
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            2 * PRECONDITIONER_DIM,
+            9,
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            1,
+        )
+        total_steps = 5
+        global_train_data = generate_global_train_data(
+            num_steps=total_steps,
+            world_size=dist.get_world_size(),
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooFSDPDistributorTest._shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
@@ -168,6 +195,9 @@ def test_fsdp_shampoo_against_default_shampoo(self) -> None:
                 post_model_decoration=partial(FSDP1, use_orig_params=True),
                 distributed_config=fsdp_config,
             ),
+            total_steps=total_steps,
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @skip_if_lt_x_gpu(2)
@@ -225,7 +255,20 @@ def test_fsdp_distributed_config_against_hsdp_distributed_config_bitwise_identic
             device_mesh=init_device_mesh("cuda", (1, self.world_size)),
             communicate_params=communicate_params,
         )
-
+        model_linear_layers_dims = (
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            2 * PRECONDITIONER_DIM,
+            9,
+            4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,
+            1,
+        )
+        total_steps = 100
+        global_train_data = generate_global_train_data(
+            num_steps=total_steps,
+            world_size=dist.get_world_size(),
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooFSDPDistributorTest._shampoo_optim_factory(
                 distributed_config=fsdp_config,
@@ -248,9 +291,11 @@ def test_fsdp_distributed_config_against_hsdp_distributed_config_bitwise_identic
                 ),
                 distributed_config=hsdp_config,
             ),
-            total_steps=100,
+            total_steps=total_steps,
             rtol=0.0,
             atol=0.0,
+            control_train_data=global_train_data[:, dist.get_rank()],
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
 
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_utils_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_utils_test.py
index 55cc587..4646570 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_utils_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_fsdp_utils_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 from collections.abc import Callable, Iterator
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_distributor_test.py
index b3fbe0c..768ea6f 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Callable
 from functools import partial
@@ -36,6 +34,7 @@
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
 from torch import distributed as dist, nn
@@ -140,6 +139,12 @@ def test_all_ranks_with_no_grads(self) -> None:
         fully_shard_config = FullyShardDistributedConfig()
 
         steps_without_gradients = 2
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=self.world_size,
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch.object(torch.Tensor, "backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
@@ -151,7 +156,7 @@ def test_all_ranks_with_no_grads(self) -> None:
                     ShampooFullyShardDistributorTest._construct_model,
                     post_model_decoration=partial(fully_shard),
                 ),
-                num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -161,6 +166,13 @@ def test_all_ranks_with_no_grads(self) -> None:
     @skip_if_lt_x_gpu(2)
     def test_fully_shard_shampoo_against_default_shampoo(self) -> None:
         fully_shard_config = FullyShardDistributedConfig()
+        total_steps = 5
+        global_train_data = generate_global_train_data(
+            num_steps=total_steps,
+            world_size=self.world_size,
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooFullyShardDistributorTest._shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
@@ -173,6 +185,9 @@ def test_fully_shard_shampoo_against_default_shampoo(self) -> None:
                 ShampooFullyShardDistributorTest._construct_model,
                 post_model_decoration=partial(fully_shard),
             ),
+            total_steps=total_steps,
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -189,6 +204,13 @@ def test_hybrid_shard_distributed_config_against_fully_shard_distributed_config_
             device_mesh=mesh_2d, communicate_params=communicate_params
         )
 
+        total_steps = 100
+        global_train_data = generate_global_train_data(
+            num_steps=total_steps,
+            world_size=self.world_size,
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooFullyShardDistributorTest._shampoo_optim_factory(
                 distributed_config=fully_shard_config
@@ -204,14 +226,22 @@ def test_hybrid_shard_distributed_config_against_fully_shard_distributed_config_
                 ShampooFullyShardDistributorTest._construct_model,
                 post_model_decoration=partial(fully_shard, mesh=mesh_2d),
             ),
-            total_steps=100,
+            total_steps=total_steps,
             rtol=0.0,
             atol=0.0,
+            control_train_data=global_train_data[:, dist.get_rank()],
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
     @skip_if_lt_x_gpu(2)
     def test_fully_shard_shampoo_block_index(self) -> None:
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         model, _, _, _, optimizer = train_model(
             optim_factory=ShampooFullyShardDistributorTest._shampoo_optim_factory(
                 distributed_config=FullyShardDistributedConfig()
@@ -220,6 +250,7 @@ def test_fully_shard_shampoo_block_index(self) -> None:
                 ShampooFullyShardDistributorTest._construct_model,
                 post_model_decoration=partial(fully_shard),
             ),
+            train_data=global_train_data[:, dist.get_rank()],
         )
         assert isinstance(model, nn.Module)
         assert isinstance(optimizer, DistributedShampoo)
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_lossless_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_lossless_distributor_test.py
index 52fd711..0e034ee 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_lossless_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_fully_shard_lossless_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Callable
 from functools import partial
@@ -22,26 +20,26 @@
 )
 from distributed_shampoo.distributor.gpu_tests.distributor_test_utils import (
     DistributorOnEmptyParamTest,
+    PARAM_GROUP_TEST_DEAD_DIMS,
+    PARAM_GROUP_TEST_MODEL_DIMS,
+    shampoo_optim_factory,
 )
 from distributed_shampoo.distributor.shampoo_block_info import BlockInfo
 from distributed_shampoo.shampoo_types import (
-    AdaGradPreconditionerConfig,
     DefaultSingleDeviceDistributedConfig,
     FSDPParamAssignmentStrategy,
     FullyShardDistributedConfig,
-    GeneralizedPrimalAveragingConfig,
-    SingleDeviceDistributedConfig,
-    WeightDecayType,
+    ShampooPT2CompileConfig,
 )
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
-from torch import nn
+from torch import distributed as dist, nn
 from torch.distributed.fsdp import FSDPModule, fully_shard
 from torch.distributed.tensor import distribute_tensor, init_device_mesh, Shard
-from torch.optim.optimizer import ParamsT
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -116,25 +114,6 @@ def _construct_model(
             fill=0.1,
         )
 
-    @staticmethod
-    def _shampoo_optim_factory(
-        distributed_config: FullyShardDistributedConfig | SingleDeviceDistributedConfig,
-    ) -> Callable[[ParamsT], torch.optim.Optimizer]:
-        return partial(
-            DistributedShampoo,
-            lr=0.001,
-            betas=(0.9, 1.0),
-            epsilon=1e-8,
-            weight_decay=0.0,
-            max_preconditioner_dim=PRECONDITIONER_DIM,
-            precondition_frequency=1,
-            start_preconditioning_step=2,
-            weight_decay_type=WeightDecayType.DECOUPLED,
-            grafting_config=AdaGradPreconditionerConfig(epsilon=1e-8),
-            distributed_config=distributed_config,
-            iterate_averaging_config=GeneralizedPrimalAveragingConfig(),
-        )
-
     @with_comms
     @skip_if_lt_x_gpu(2)
     def test_init_assign_full_parameters(self) -> None:
@@ -195,11 +174,17 @@ def test_all_ranks_with_no_grads(
         )
 
         steps_without_gradients = 2
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=self.world_size,
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch("torch.Tensor.backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
             train_model(
-                optim_factory=ShampooFullyShardLosslessDistributorTest._shampoo_optim_factory(
+                optim_factory=shampoo_optim_factory(
                     distributed_config=fully_shard_config,
                 ),
                 model_factory=partial(
@@ -208,7 +193,7 @@ def test_all_ranks_with_no_grads(
                     model_dead_layers_dims=model_dead_layers_dims,
                     post_model_decoration=partial(fully_shard),
                 ),
-                num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -239,25 +224,39 @@ def test_fully_shard_shampoo_against_default_shampoo(
             model_linear_layers_dims=model_linear_layers_dims,
             model_dead_layers_dims=model_dead_layers_dims,
         )
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
-            control_optim_factory=ShampooFullyShardLosslessDistributorTest._shampoo_optim_factory(
+            control_optim_factory=shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
             ),
             control_model_factory=control_model_factory,
-            experimental_optim_factory=ShampooFullyShardLosslessDistributorTest._shampoo_optim_factory(
+            experimental_optim_factory=shampoo_optim_factory(
                 distributed_config=fully_shard_config,
             ),
             experimental_model_factory=partial(
                 control_model_factory,
                 post_model_decoration=partial(fully_shard),
             ),
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
     @skip_if_lt_x_gpu(2)
     def test_fully_shard_lossless_shampoo_block_index(self) -> None:
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(TEST_MODEL_LAYER_DIMS[0][0],),
+            device=torch.device("cuda"),
+        )
         model, _, _, _, optimizer = train_model(
-            optim_factory=ShampooFullyShardLosslessDistributorTest._shampoo_optim_factory(
+            optim_factory=shampoo_optim_factory(
                 distributed_config=FullyShardDistributedConfig(
                     param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE
                 )
@@ -268,6 +267,7 @@ def test_fully_shard_lossless_shampoo_block_index(self) -> None:
                 model_dead_layers_dims=DEAD_MODEL_LAYER_DIMS[0],
                 post_model_decoration=partial(fully_shard),
             ),
+            train_data=global_train_data[:, dist.get_rank()],
         )
         assert isinstance(model, nn.Module)
         assert isinstance(optimizer, DistributedShampoo)
@@ -422,3 +422,95 @@ def _expected_local_masked_block_grads(self) -> tuple[torch.Tensor, ...]:
     @with_comms
     def test_merge_and_block_gradients(self) -> None:  # type: ignore[override]
         DistributorOnEmptyParamTest.Interface.test_merge_and_block_gradients(self)
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+@instantiate_parametrized_tests
+class ShampooFullyShardParamGroupTest(DTensorTestBase):
+    """Tests for param group splitting with num_sub_groups > 1."""
+
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    @parametrize("num_sub_groups", (2, 3, 4))
+    @parametrize(
+        "shampoo_pt2_compile_config",
+        (ShampooPT2CompileConfig(), None),
+    )
+    def test_fully_shard_shampoo_split_matches_unsplit(
+        self,
+        num_sub_groups: int,
+        shampoo_pt2_compile_config: ShampooPT2CompileConfig | None,
+    ) -> None:
+        """Verify split FSDP (num_sub_groups > 1) matches unsplit single-group."""
+        control_config = FullyShardDistributedConfig(
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=1,
+        )
+        experimental_config = FullyShardDistributedConfig(
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=num_sub_groups,
+        )
+        model_factory = partial(
+            construct_training_problem,
+            model_linear_layers_dims=PARAM_GROUP_TEST_MODEL_DIMS,
+            model_dead_layers_dims=PARAM_GROUP_TEST_DEAD_DIMS,
+            enable_learnable_scalar=False,
+            device=torch.device("cuda"),
+            fill=0.1,
+            post_model_decoration=partial(fully_shard),
+        )
+        compare_two_optimizers_models_devices_on_weight_and_loss(
+            control_optim_factory=shampoo_optim_factory(
+                distributed_config=control_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            control_model_factory=model_factory,
+            experimental_optim_factory=shampoo_optim_factory(
+                distributed_config=experimental_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            experimental_model_factory=model_factory,
+        )
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    @parametrize(
+        "shampoo_pt2_compile_config",
+        (ShampooPT2CompileConfig(), None),
+    )
+    def test_fully_shard_shampoo_with_num_sub_groups_against_default_shampoo(
+        self,
+        shampoo_pt2_compile_config: ShampooPT2CompileConfig | None,
+    ) -> None:
+        """Verify split FSDP (num_sub_groups=2) matches single-device default Shampoo."""
+        fully_shard_config = FullyShardDistributedConfig(
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=2,
+        )
+        control_model_factory = partial(
+            construct_training_problem,
+            model_linear_layers_dims=PARAM_GROUP_TEST_MODEL_DIMS,
+            model_dead_layers_dims=PARAM_GROUP_TEST_DEAD_DIMS,
+            enable_learnable_scalar=False,
+            device=torch.device("cuda"),
+            fill=0.1,
+        )
+        compare_two_optimizers_models_devices_on_weight_and_loss(
+            control_optim_factory=shampoo_optim_factory(
+                distributed_config=DefaultSingleDeviceDistributedConfig,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            control_model_factory=control_model_factory,
+            experimental_optim_factory=shampoo_optim_factory(
+                distributed_config=fully_shard_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            experimental_model_factory=partial(
+                control_model_factory,
+                post_model_decoration=partial(fully_shard),
+            ),
+        )
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_hsdp_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_hsdp_distributor_test.py
index 1678109..cc632c4 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_hsdp_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_hsdp_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 from collections.abc import Callable
@@ -38,6 +36,7 @@
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
 from torch import distributed as dist, nn
@@ -154,6 +153,12 @@ def test_hsdp_shampoo_against_default_shampoo(
             communicate_params=communicate_params,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooHSDPDistributorTest._shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
@@ -172,6 +177,8 @@ def test_hsdp_shampoo_against_default_shampoo(
                 ),
                 distributed_config=hsdp_config,
             ),
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @skip_if_lt_x_gpu(4)
@@ -182,6 +189,12 @@ def test_hsdp_shampoo_block_index(self) -> None:
             device_mesh=mesh_2d,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         model, _, _, _, optimizer = train_model(
             optim_factory=ShampooHSDPDistributorTest._shampoo_optim_factory(
                 hsdp_config
@@ -196,6 +209,7 @@ def test_hsdp_shampoo_block_index(self) -> None:
                 ),
                 distributed_config=hsdp_config,
             ),
+            train_data=global_train_data[:, dist.get_rank()],
         )
         assert isinstance(optimizer, DistributedShampoo)
         osd_state = optimizer.state_dict()["state"]
@@ -243,6 +257,12 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
         )
 
         steps_without_gradients = 2
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch.object(torch.Tensor, "backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
@@ -260,7 +280,7 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
                     ),
                     distributed_config=hsdp_config,
                 ),
-                num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_distributor_test.py
index 4fdcfd3..129f070 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import math
 import re
 import unittest
@@ -40,12 +38,15 @@
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
 from torch import distributed as dist, nn
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import FSDPModule, fully_shard
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.optimizer import ParamsT
+from torch.testing._comparison import default_tolerances
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -176,6 +177,12 @@ def test_hybrid_shard_shampoo_against_default_shampoo(
             communicate_params=communicate_params,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
@@ -190,6 +197,8 @@ def test_hybrid_shard_shampoo_against_default_shampoo(
                     fully_shard, mesh=hybrid_shard_config.device_mesh
                 ),
             ),
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -224,6 +233,12 @@ def test_hybrid_shampoo_n_by_one_mesh_against_default_shampoo(
             communicate_params=communicate_params,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
@@ -238,6 +253,8 @@ def test_hybrid_shampoo_n_by_one_mesh_against_default_shampoo(
                     fully_shard, mesh=hybrid_shard_config.device_mesh
                 ),
             ),
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -246,12 +263,26 @@ def test_hybrid_shampoo_n_by_one_mesh_against_default_shampoo(
         "start_preconditioning_step", (2, math.inf)
     )  # math.inf here is to test the grafting similarities between HSDP2 and DDP
     @parametrize(
-        "communication_dtype, communicate_params",
+        "communication_dtype, communicate_params, rtol, atol",
         (
-            (torch.float32, False),
-            (torch.float32, True),
-            (torch.float16, False),
-            (torch.bfloat16, False),
+            # float32 collectives: bitwise identical expected because both DDP
+            # and HybridShard(n,1) perform NCCL all-reduce on the same inputs
+            # and topology — no dtype casting means no rounding divergence.
+            (torch.float32, False, 0.0, 0.0),
+            (torch.float32, True, 0.0, 0.0),
+            # Reduced-precision collectives: 2x bfloat16 default tolerances
+            # (rtol=0.032, atol=2e-5) to account for error amplification
+            # through Shampoo preconditioner (eigendecomposition / root inverse).
+            (
+                torch.float16,
+                False,
+                *[2 * tol for tol in default_tolerances(torch.bfloat16)],
+            ),
+            (
+                torch.bfloat16,
+                False,
+                *[2 * tol for tol in default_tolerances(torch.bfloat16)],
+            ),
         ),
     )
     @parametrize("num_trainers_per_group", (-1, 1, 2, 4))
@@ -261,6 +292,8 @@ def test_hybrid_shampoo_n_by_one_mesh_against_ddp_shampoo(
         communication_dtype: torch.dtype,
         communicate_params: bool,
         start_preconditioning_step: int,
+        rtol: float,
+        atol: float,
     ) -> None:
         """
         Testing the correctness of hybrid shard Shampoo distributor of (n, 1) mesh
@@ -280,12 +313,24 @@ def test_hybrid_shampoo_n_by_one_mesh_against_ddp_shampoo(
             communicate_params=communicate_params,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=ddp_config,
                 start_preconditioning_step=start_preconditioning_step,
             ),
-            control_model_factory=ShampooHybridShardDistributorTest._construct_model,
+            control_model_factory=partial(
+                ShampooHybridShardDistributorTest._construct_model,
+                # DDP wrapping is needed so gradients are all-reduced across ranks.
+                # find_unused_parameters=True is required because the model's dead
+                # layers don't participate in forward() and never receive gradients.
+                post_model_decoration=partial(DDP, find_unused_parameters=True),
+            ),
             experimental_optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=hybrid_shard_config,
                 start_preconditioning_step=start_preconditioning_step,
@@ -296,19 +341,35 @@ def test_hybrid_shampoo_n_by_one_mesh_against_ddp_shampoo(
                     fully_shard, mesh=hybrid_shard_config.device_mesh
                 ),
             ),
-            rtol=0.0,
-            atol=0.0,
+            rtol=rtol,
+            atol=atol,
+            control_train_data=global_train_data[:, dist.get_rank()],
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
     @skip_if_lt_x_gpu(4)
     @parametrize(
-        "communication_dtype, communicate_params",
+        "communication_dtype, communicate_params, rtol, atol",
         (
-            (torch.float32, False),
-            (torch.float32, True),
-            (torch.float16, False),
-            (torch.bfloat16, False),
+            # float32 collectives: bitwise identical expected because both DDP
+            # and HybridShard(n,1) perform NCCL all-reduce on the same inputs
+            # and topology — no dtype casting means no rounding divergence.
+            (torch.float32, False, 0.0, 0.0),
+            (torch.float32, True, 0.0, 0.0),
+            # Reduced-precision collectives: 2x bfloat16 default tolerances
+            # (rtol=0.032, atol=2e-5) to account for error amplification
+            # through Shampoo preconditioner (eigendecomposition / root inverse).
+            (
+                torch.float16,
+                False,
+                *[2 * tol for tol in default_tolerances(torch.bfloat16)],
+            ),
+            (
+                torch.bfloat16,
+                False,
+                *[2 * tol for tol in default_tolerances(torch.bfloat16)],
+            ),
         ),
     )
     @parametrize("num_trainers_per_group", (-1, 1, 2))
@@ -317,6 +378,8 @@ def test_hybrid_shard_distributed_config_against_fully_shard_distributed_config(
         num_trainers_per_group: int,
         communication_dtype: torch.dtype,
         communicate_params: bool,
+        rtol: float,
+        atol: float,
     ) -> None:
         """
         Testing the correctness of hybrid shard shampoo distributor by comparing it with
@@ -333,6 +396,12 @@ def test_hybrid_shard_distributed_config_against_fully_shard_distributed_config(
             communicate_params=communicate_params,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
             control_optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=fully_shard_config
@@ -348,6 +417,10 @@ def test_hybrid_shard_distributed_config_against_fully_shard_distributed_config(
                 ShampooHybridShardDistributorTest._construct_model,
                 post_model_decoration=partial(fully_shard, mesh=mesh_2d),
             ),
+            rtol=rtol,
+            atol=atol,
+            control_train_data=global_train_data[:, dist.get_rank()],
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -362,6 +435,12 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
         )
 
         steps_without_gradients = 2
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch.object(torch.Tensor, "backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
@@ -375,7 +454,7 @@ def test_all_ranks_with_no_grads(self, communicate_params: bool) -> None:
                         fully_shard, mesh=hybrid_shard_config.device_mesh
                     ),
                 ),
-                num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -387,6 +466,12 @@ def test_hybrid_shard_shampoo_block_index(self) -> None:
         mesh_2d = init_device_mesh(
             "cuda", (2, 2), mesh_dim_names=("replicate", "shard")
         )
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=dist.get_world_size(),
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )
         model, _, _, _, optimizer = train_model(
             optim_factory=ShampooHybridShardDistributorTest._shampoo_optim_factory(
                 distributed_config=HybridShardDistributedConfig(device_mesh=mesh_2d)
@@ -395,6 +480,7 @@ def test_hybrid_shard_shampoo_block_index(self) -> None:
                 ShampooHybridShardDistributorTest._construct_model,
                 post_model_decoration=partial(fully_shard, mesh=mesh_2d),
             ),
+            train_data=global_train_data[:, dist.get_rank()],
         )
         assert isinstance(model, nn.Module)
         assert isinstance(optimizer, DistributedShampoo)
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_lossless_distributor_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_lossless_distributor_test.py
index b9f78e4..41575da 100644
--- a/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_lossless_distributor_test.py
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_hybrid_shard_lossless_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Callable
 from functools import partial
@@ -22,28 +20,27 @@
 )
 from distributed_shampoo.distributor.gpu_tests.distributor_test_utils import (
     DistributorOnEmptyParamTest,
+    HYBRID_SHARD_PARAM_GROUP_TEST_MODEL_DIMS,
+    PARAM_GROUP_TEST_DEAD_DIMS,
+    shampoo_optim_factory,
 )
 from distributed_shampoo.distributor.shampoo_block_info import DTensorBlockInfo
 from distributed_shampoo.shampoo_types import (
-    AdaGradPreconditionerConfig,
-    DDPDistributedConfig,
     DefaultSingleDeviceDistributedConfig,
     FSDPParamAssignmentStrategy,
     FullyShardDistributedConfig,
-    GeneralizedPrimalAveragingConfig,
     HybridShardDistributedConfig,
-    SingleDeviceDistributedConfig,
-    WeightDecayType,
+    ShampooPT2CompileConfig,
 )
 from distributed_shampoo.tests.shampoo_test_utils import (
     compare_two_optimizers_models_devices_on_weight_and_loss,
     construct_training_problem,
+    generate_global_train_data,
     train_model,
 )
 from torch import distributed as dist, nn
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import FSDPModule, fully_shard
-from torch.optim.optimizer import ParamsT
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -118,29 +115,6 @@ def _construct_model(
             fill=0.1,
         )
 
-    @staticmethod
-    def _shampoo_optim_factory(
-        distributed_config: DDPDistributedConfig
-        | FullyShardDistributedConfig
-        | HybridShardDistributedConfig
-        | SingleDeviceDistributedConfig,
-        start_preconditioning_step: int = 2,
-    ) -> Callable[[ParamsT], torch.optim.Optimizer]:
-        return partial(
-            DistributedShampoo,
-            lr=0.001,
-            betas=(0.9, 1.0),
-            epsilon=1e-8,
-            weight_decay=0.0,
-            max_preconditioner_dim=PRECONDITIONER_DIM,
-            precondition_frequency=1,
-            start_preconditioning_step=start_preconditioning_step,
-            weight_decay_type=WeightDecayType.DECOUPLED,
-            grafting_config=AdaGradPreconditionerConfig(epsilon=1e-8),
-            distributed_config=distributed_config,
-            iterate_averaging_config=GeneralizedPrimalAveragingConfig(),
-        )
-
     @with_comms
     @skip_if_lt_x_gpu(4)
     @parametrize(
@@ -181,8 +155,14 @@ def test_hybrid_shard_shampoo_against_default_shampoo(
             param_assignment_strategy=param_assignment_strategy,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
-            control_optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+            control_optim_factory=shampoo_optim_factory(
                 distributed_config=DefaultSingleDeviceDistributedConfig,
             ),
             control_model_factory=partial(
@@ -190,7 +170,7 @@ def test_hybrid_shard_shampoo_against_default_shampoo(
                 model_linear_layers_dims=model_linear_layers_dims,
                 model_dead_layers_dims=model_dead_layers_dims,
             ),
-            experimental_optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+            experimental_optim_factory=shampoo_optim_factory(
                 distributed_config=hybrid_shard_config,
             ),
             experimental_model_factory=partial(
@@ -201,6 +181,8 @@ def test_hybrid_shard_shampoo_against_default_shampoo(
                     fully_shard, mesh=hybrid_shard_config.device_mesh
                 ),
             ),
+            control_train_data=global_train_data,
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -249,8 +231,14 @@ def test_hybrid_shard_shampoo_config_against_fully_shard_shampoo_config(
             param_assignment_strategy=param_assignment_strategy,
         )
 
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         compare_two_optimizers_models_devices_on_weight_and_loss(
-            control_optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+            control_optim_factory=shampoo_optim_factory(
                 distributed_config=fully_shard_config
             ),
             control_model_factory=partial(
@@ -259,7 +247,7 @@ def test_hybrid_shard_shampoo_config_against_fully_shard_shampoo_config(
                 model_dead_layers_dims=(),
                 post_model_decoration=partial(fully_shard),
             ),
-            experimental_optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+            experimental_optim_factory=shampoo_optim_factory(
                 distributed_config=hybrid_shard_config
             ),
             experimental_model_factory=partial(
@@ -268,6 +256,8 @@ def test_hybrid_shard_shampoo_config_against_fully_shard_shampoo_config(
                 model_dead_layers_dims=(),
                 post_model_decoration=partial(fully_shard, mesh=mesh_2d),
             ),
+            control_train_data=global_train_data[:, dist.get_rank()],
+            experimental_train_data=global_train_data[:, dist.get_rank()],
         )
 
     @with_comms
@@ -290,11 +280,17 @@ def test_all_ranks_with_no_grads(
         )
 
         steps_without_gradients = 2
+        global_train_data = generate_global_train_data(
+            num_steps=steps_without_gradients,
+            world_size=self.world_size,
+            data_shape=(model_linear_layers_dims[0],),
+            device=torch.device("cuda"),
+        )
         with unittest.mock.patch("torch.Tensor.backward") as mock_backward:
             # By mocking the backward() method, we're intercepting gradient calculation.
             # This effectively simulates running forward passes without computing gradients.
             train_model(
-                optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+                optim_factory=shampoo_optim_factory(
                     distributed_config=hybrid_shard_config
                 ),
                 model_factory=partial(
@@ -305,7 +301,7 @@ def test_all_ranks_with_no_grads(
                         fully_shard, mesh=hybrid_shard_config.device_mesh
                     ),
                 ),
-                num_steps=steps_without_gradients,
+                train_data=global_train_data[:, dist.get_rank()],
             )
 
         # Verify that the backward() method was called the expected number of times and the training loop completed successfully.
@@ -317,8 +313,14 @@ def test_hybrid_shard_lossless_shampoo_block_index(self) -> None:
         mesh_2d = init_device_mesh(
             "cuda", (2, 2), mesh_dim_names=("replicate", "shard")
         )
+        global_train_data = generate_global_train_data(
+            num_steps=5,
+            world_size=self.world_size,
+            data_shape=(TEST_MODEL_LAYER_DIMS[0][0],),
+            device=torch.device("cuda"),
+        )
         model, _, _, _, optimizer = train_model(
-            optim_factory=ShampooHybridShardLosslessDistributorTest._shampoo_optim_factory(
+            optim_factory=shampoo_optim_factory(
                 distributed_config=HybridShardDistributedConfig(
                     device_mesh=mesh_2d,
                     param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE,
@@ -330,6 +332,7 @@ def test_hybrid_shard_lossless_shampoo_block_index(self) -> None:
                 model_dead_layers_dims=DEAD_MODEL_LAYER_DIMS[0],
                 post_model_decoration=partial(fully_shard, mesh=mesh_2d),
             ),
+            train_data=global_train_data[:, dist.get_rank()],
         )
         assert isinstance(model, nn.Module)
         assert isinstance(optimizer, DistributedShampoo)
@@ -500,3 +503,108 @@ def _expected_local_masked_block_grads(self) -> tuple[torch.Tensor, ...]:
     @with_comms
     def test_merge_and_block_gradients(self) -> None:  # type: ignore[override]
         DistributorOnEmptyParamTest.Interface.test_merge_and_block_gradients(self)
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+@instantiate_parametrized_tests
+class ShampooHybridShardParamGroupTest(DTensorTestBase):
+    """Tests for param group splitting with HybridShardDistributedConfig and num_sub_groups > 1."""
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize("num_sub_groups", (2, 3, 4))
+    @parametrize(
+        "shampoo_pt2_compile_config",
+        (ShampooPT2CompileConfig(), None),
+    )
+    def test_hybrid_shard_shampoo_split_matches_unsplit(
+        self,
+        num_sub_groups: int,
+        shampoo_pt2_compile_config: ShampooPT2CompileConfig | None,
+    ) -> None:
+        """Verify split HSDP (num_sub_groups > 1) matches unsplit single-group."""
+        device_mesh = init_device_mesh(
+            "cuda", (2, 2), mesh_dim_names=("replicate", "shard")
+        )
+        control_config = HybridShardDistributedConfig(
+            device_mesh=device_mesh,
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=1,
+        )
+        experimental_config = HybridShardDistributedConfig(
+            device_mesh=device_mesh,
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=num_sub_groups,
+        )
+        model_factory = partial(
+            construct_training_problem,
+            model_linear_layers_dims=HYBRID_SHARD_PARAM_GROUP_TEST_MODEL_DIMS,
+            model_dead_layers_dims=PARAM_GROUP_TEST_DEAD_DIMS,
+            enable_learnable_scalar=False,
+            device=torch.device("cuda"),
+            fill=0.1,
+            post_model_decoration=partial(fully_shard, mesh=device_mesh),
+        )
+        compare_two_optimizers_models_devices_on_weight_and_loss(
+            control_optim_factory=shampoo_optim_factory(
+                distributed_config=control_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            control_model_factory=model_factory,
+            experimental_optim_factory=shampoo_optim_factory(
+                distributed_config=experimental_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            experimental_model_factory=model_factory,
+        )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize(
+        "shampoo_pt2_compile_config",
+        (ShampooPT2CompileConfig(), None),
+    )
+    def test_hybrid_shard_shampoo_with_num_sub_groups_against_default_shampoo(
+        self,
+        shampoo_pt2_compile_config: ShampooPT2CompileConfig | None,
+    ) -> None:
+        """Verify split HSDP (num_sub_groups=2) matches single-device default Shampoo."""
+        device_mesh = init_device_mesh(
+            "cuda", (2, 2), mesh_dim_names=("replicate", "shard")
+        )
+        hybrid_shard_config = HybridShardDistributedConfig(
+            device_mesh=device_mesh,
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=2,
+        )
+        compare_two_optimizers_models_devices_on_weight_and_loss(
+            control_optim_factory=shampoo_optim_factory(
+                distributed_config=DefaultSingleDeviceDistributedConfig,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            control_model_factory=partial(
+                construct_training_problem,
+                model_linear_layers_dims=HYBRID_SHARD_PARAM_GROUP_TEST_MODEL_DIMS,
+                model_dead_layers_dims=PARAM_GROUP_TEST_DEAD_DIMS,
+                enable_learnable_scalar=False,
+                device=torch.device("cuda"),
+                fill=0.1,
+            ),
+            experimental_optim_factory=shampoo_optim_factory(
+                distributed_config=hybrid_shard_config,
+                shampoo_pt2_compile_config=shampoo_pt2_compile_config,
+            ),
+            experimental_model_factory=partial(
+                construct_training_problem,
+                model_linear_layers_dims=HYBRID_SHARD_PARAM_GROUP_TEST_MODEL_DIMS,
+                model_dead_layers_dims=PARAM_GROUP_TEST_DEAD_DIMS,
+                enable_learnable_scalar=False,
+                device=torch.device("cuda"),
+                fill=0.1,
+                post_model_decoration=partial(fully_shard, mesh=device_mesh),
+            ),
+        )
diff --git a/distributed_shampoo/distributor/gpu_tests/shampoo_recompilation_test.py b/distributed_shampoo/distributor/gpu_tests/shampoo_recompilation_test.py
new file mode 100644
index 0000000..afc3f7d
--- /dev/null
+++ b/distributed_shampoo/distributor/gpu_tests/shampoo_recompilation_test.py
@@ -0,0 +1,339 @@
+#!/usr/bin/env python3
+
+
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+Regression tests asserting that DistributedShampoo's torch.compile recompilation count
+is bounded.
+
+With ShampooPT2CompileConfig enabled, the only Python-bool input to the per-group step
+function that varies across steps AND enters a compiled region is
+``perform_amortized_computation``. (``use_grafting_method`` also varies, but it is
+consumed entirely inside ``_precondition_and_grafting`` which is decorated with
+``@torch.compiler.disable`` — see distributed_shampoo.py — so it cannot specialize a
+compiled graph.) Across a run with start_preconditioning_step=S and
+precondition_frequency=2, dynamo should specialize on exactly two values of
+``perform_amortized_computation`` and therefore trigger compilation at exactly two
+steps:
+  - step 1 (perform_amortized=False; first call, fresh compile)
+  - step S (perform_amortized=True; first amortized inverse-root recomputation, fresh
+    compile)
+Subsequent steps reuse the cached graph for whichever value of
+``perform_amortized_computation`` they hit and must not trigger any further compilation.
+
+A new compilation at any other step indicates a regression where new dynamic state has
+leaked into the compiled region (e.g., a new Python-int input that varies across steps,
+a new control-flow branch keyed on a Python int, a new dynamic-shape guard).
+
+Scope: covers the actively used distributors (single-device, DDP, FullyShard family,
+HybridShard family). Each multi-proc distributor lives in its own test class so that
+MultiProcessTestCase spawns fresh subprocesses per variant; sharing one parametrized
+class across all five was observed to flake on NCCL/dynamo state cleanup. The legacy
+FSDPDistributor (FSDP1) and HSDPDistributor (FSDP1-based HSDP) are out of scope here
+since they require the FSDPTest harness plus compile_fsdp_parameter_metadata; HSDP
+topology with the modern fully_shard backend IS covered via the hybrid_shard classes.
+
+Numerical equivalence between the compiled and eager paths is covered separately in
+dev/gpu_tests/shampoo_pt2_test.py; this file owns the orthogonal compile-count invariant.
+"""
+
+import unittest
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from functools import partial
+
+import torch
+import torch._dynamo
+from distributed_shampoo.distributed_shampoo import DistributedShampoo
+from distributed_shampoo.shampoo_types import (
+    AdaGradPreconditionerConfig,
+    DDPDistributedConfig,
+    DefaultSingleDeviceDistributedConfig,
+    DistributedConfig,
+    FSDPParamAssignmentStrategy,
+    FullyShardDistributedConfig,
+    HybridShardDistributedConfig,
+    ShampooPT2CompileConfig,
+    WeightDecayType,
+)
+from distributed_shampoo.tests.shampoo_test_utils import (
+    construct_training_problem,
+    generate_global_train_data,
+)
+from torch import distributed as dist, nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import fully_shard
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim.optimizer import ParamsT
+from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
+from torch.testing._internal.distributed._tensor.common_dtensor import (
+    DTensorTestBase,
+    with_comms,
+)
+
+
+PRECONDITIONER_DIM = 4
+PRECONDITION_FREQUENCY = 2
+START_PRECONDITIONING_STEP = 2
+NUM_STEPS = 6  # Run past S+1 to verify steady state has no further recompiles.
+
+# 1-indexed step numbers at which dynamo is expected to compile a new graph.
+# Assumes START_PRECONDITIONING_STEP >= 2 so the two compile points are distinct.
+# NOTE: dynamo emits MULTIPLE sub-graphs per top-level _per_group_step_impl call (due to
+# graph breaks at @torch.compiler.disable boundaries and similar). This test only asserts
+# WHICH steps trigger compilation, not the per-step graph COUNT — so a regression that
+# inflates the number of sub-graphs within an already-expected step (e.g., an
+# accidentally-introduced graph break) will not be caught here. Pinning per-distributor
+# graph-count baselines is a follow-up.
+EXPECTED_COMPILE_STEPS: frozenset[int] = frozenset({1, START_PRECONDITIONING_STEP})
+
+
+@contextmanager
+def _fresh_dynamo_state() -> Iterator[None]:
+    torch._dynamo.reset()
+    torch._dynamo.utils.counters.clear()
+    try:
+        yield
+    finally:
+        torch._dynamo.reset()
+        torch._dynamo.utils.counters.clear()
+
+
+def _shampoo_optim_factory(
+    distributed_config: DistributedConfig,
+) -> Callable[[ParamsT], torch.optim.Optimizer]:
+    return partial(
+        DistributedShampoo,
+        lr=0.001,
+        betas=(0.9, 1.0),
+        epsilon=1e-8,
+        weight_decay=0.0,
+        max_preconditioner_dim=PRECONDITIONER_DIM,
+        precondition_frequency=PRECONDITION_FREQUENCY,
+        start_preconditioning_step=START_PRECONDITIONING_STEP,
+        weight_decay_type=WeightDecayType.DECOUPLED,
+        grafting_config=AdaGradPreconditionerConfig(epsilon=1e-8),
+        distributed_config=distributed_config,
+        shampoo_pt2_compile_config=ShampooPT2CompileConfig(),
+    )
+
+
+def _run_and_assert_recompiles(
+    test_case: unittest.TestCase,
+    optimizer: torch.optim.Optimizer,
+    model: nn.Module,
+    loss_fn: nn.Module,
+    train_data: torch.Tensor,
+    target: torch.Tensor,
+) -> None:
+    """Run NUM_STEPS optimizer.step() calls; assert the set of steps that triggered any
+    dynamo compilation equals EXPECTED_COMPILE_STEPS.
+
+    Mechanism: snapshot ``torch._dynamo.utils.counters["stats"]["unique_graphs"]`` before
+    each step; any positive delta after step k means dynamo emitted at least one new
+    graph at step k. Catches regressions that compile at unexpected steps OR fail to
+    compile at expected steps.
+    """
+
+    def unique_graphs() -> int:
+        return torch._dynamo.utils.counters["stats"]["unique_graphs"]
+
+    per_step_deltas: dict[int, int] = {}
+    last_count = unique_graphs()
+
+    for step in range(1, NUM_STEPS + 1):
+        optimizer.zero_grad()
+        objective = loss_fn(model(train_data[step - 1]), target)
+        objective.backward()
+        optimizer.step()
+
+        current_count = unique_graphs()
+        per_step_deltas[step] = current_count - last_count
+        last_count = current_count
+
+    actual_compile_steps = {step for step, d in per_step_deltas.items() if d > 0}
+
+    test_case.assertEqual(
+        actual_compile_steps,
+        set(EXPECTED_COMPILE_STEPS),
+        msg=(
+            "Dynamo compilation steps mismatch.\n"
+            f"  Expected (compile only at): {sorted(EXPECTED_COMPILE_STEPS)}\n"
+            f"  Actual   (compiled at):     {sorted(actual_compile_steps)}\n"
+            f"  Per-step deltas:            {per_step_deltas}\n"
+            "Each unexpected step represents a recompilation regression — most often "
+            "caused by a new Python-int/bool input to _per_group_step_impl that varies "
+            "across steps, a new dynamic-shape guard, or new control flow in the "
+            "Shampoo per-group step path."
+        ),
+    )
+
+
+def _run_distributed_recompilation_test(
+    test_case: DTensorTestBase,
+    distributed_config: DistributedConfig,
+    post_model_decoration: Callable[[nn.Module], nn.Module],
+) -> None:
+    """Shared body for every distributed-distributor recompilation test class."""
+    with _fresh_dynamo_state():
+        model, loss_fn, _, target = construct_training_problem(
+            model_linear_layers_dims=(
+                4 * PRECONDITIONER_DIM,
+                2 * PRECONDITIONER_DIM,
+                1,
+            ),
+            model_dead_layers_dims=None,
+            enable_learnable_scalar=False,
+            device=torch.device("cuda"),
+            fill=0.1,
+            post_model_decoration=post_model_decoration,
+        )
+        assert isinstance(model, nn.Module)
+        train_data = generate_global_train_data(
+            num_steps=NUM_STEPS,
+            world_size=test_case.world_size,
+            data_shape=(4 * PRECONDITIONER_DIM,),
+            device=torch.device("cuda"),
+        )[:, dist.get_rank()]
+
+        optimizer = _shampoo_optim_factory(
+            distributed_config=distributed_config,
+        )(model.parameters())
+
+        _run_and_assert_recompiles(
+            test_case, optimizer, model, loss_fn, train_data, target
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooSingleDeviceRecompilationTest(unittest.TestCase):
+    def test_recompilation_count(self) -> None:
+        with _fresh_dynamo_state():
+            model, loss_fn, _, target = construct_training_problem(
+                model_linear_layers_dims=(
+                    4 * PRECONDITIONER_DIM,
+                    2 * PRECONDITIONER_DIM,
+                    1,
+                ),
+                model_dead_layers_dims=None,
+                enable_learnable_scalar=False,
+                device=torch.device("cuda"),
+                fill=0.1,
+            )
+            assert isinstance(model, nn.Module)
+            train_data = generate_global_train_data(
+                num_steps=NUM_STEPS,
+                world_size=1,
+                data_shape=(4 * PRECONDITIONER_DIM,),
+                device=torch.device("cuda"),
+            ).squeeze(1)
+
+            optimizer = _shampoo_optim_factory(
+                distributed_config=DefaultSingleDeviceDistributedConfig,
+            )(model.parameters())
+
+            _run_and_assert_recompiles(
+                self, optimizer, model, loss_fn, train_data, target
+            )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooDDPRecompilationTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_recompilation_count(self) -> None:
+        _run_distributed_recompilation_test(
+            self,
+            distributed_config=DDPDistributedConfig(),
+            post_model_decoration=partial(
+                DDP, device_ids=[self.rank], find_unused_parameters=False
+            ),
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooFullyShardRecompilationTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_recompilation_count(self) -> None:
+        _run_distributed_recompilation_test(
+            self,
+            distributed_config=FullyShardDistributedConfig(),
+            post_model_decoration=partial(fully_shard),  # type: ignore
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooFullyShardLosslessRecompilationTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_recompilation_count(self) -> None:
+        _run_distributed_recompilation_test(
+            self,
+            distributed_config=FullyShardDistributedConfig(
+                param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE,
+            ),
+            post_model_decoration=partial(fully_shard),  # type: ignore
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooHybridShardRecompilationTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_recompilation_count(self) -> None:
+        mesh_2d = init_device_mesh(
+            "cuda",
+            (1, self.world_size),
+            mesh_dim_names=("replicate", "shard"),
+        )
+        _run_distributed_recompilation_test(
+            self,
+            distributed_config=HybridShardDistributedConfig(device_mesh=mesh_2d),
+            post_model_decoration=partial(fully_shard, mesh=mesh_2d),  # type: ignore
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+class ShampooHybridShardLosslessRecompilationTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    @with_comms
+    @skip_if_lt_x_gpu(2)
+    def test_recompilation_count(self) -> None:
+        mesh_2d = init_device_mesh(
+            "cuda",
+            (1, self.world_size),
+            mesh_dim_names=("replicate", "shard"),
+        )
+        _run_distributed_recompilation_test(
+            self,
+            distributed_config=HybridShardDistributedConfig(
+                device_mesh=mesh_2d,
+                param_assignment_strategy=FSDPParamAssignmentStrategy.REPLICATE,
+            ),
+            post_model_decoration=partial(fully_shard, mesh=mesh_2d),  # type: ignore
+        )
diff --git a/distributed_shampoo/distributor/shampoo_ddp_distributor.py b/distributed_shampoo/distributor/shampoo_ddp_distributor.py
index a282b7b..df38642 100644
--- a/distributed_shampoo/distributor/shampoo_ddp_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_ddp_distributor.py
@@ -190,6 +190,7 @@ def __init__(
             load_balancing_config=distributed_config.load_balancing_config,
         )
 
+        # pyrefly: ignore [bad-override-mutable-attribute]
         self._local_block_info_list: tuple[DTensorBlockInfo, ...] = (
             self._construct_local_block_info_list(
                 group_source_ranks=tuple(
@@ -573,7 +574,7 @@ def _allocate_zeros_distributed_tensor(
         ranks_in_group = dist.get_process_group_ranks(group=self._dist_group)
         device_mesh_2d = get_device_mesh(
             device_type=device.type,
-            mesh=tuple(batched(iterable=ranks_in_group, n=self._group_size)),
+            mesh=tuple(batched(iterable=ranks_in_group, n=self._group_size)),  # noqa: B911
             mesh_dim_names=("replicate", "shard"),
         )
         replicate_submesh = device_mesh_2d._get_all_submeshes(  # type: ignore[attr-defined]
diff --git a/distributed_shampoo/distributor/shampoo_dist_utils.py b/distributed_shampoo/distributor/shampoo_dist_utils.py
index 629e40b..1b98b05 100644
--- a/distributed_shampoo/distributor/shampoo_dist_utils.py
+++ b/distributed_shampoo/distributor/shampoo_dist_utils.py
@@ -8,9 +8,11 @@
 """
 
 from contextlib import contextmanager
-from functools import cache
+from functools import cache, partial
 from typing import Generator
 
+import torch
+from torch import distributed as dist
 from torch.autograd import profiler
 from torch.distributed.device_mesh import DeviceMesh
 
@@ -27,10 +29,34 @@ def shampoo_comm_profiler(name: str) -> Generator[None, None, None]:
             dist.all_gather_into_tensor(...)
 
     """
+    # TODO(irisz): Investigate adding CUDA NVTX ranges (torch.cuda.nvtx.range_push/pop)
+    # so annotations appear on the GPU timeline even when launched from worker threads.
     with profiler.record_function(name):
         yield
 
 
+@contextmanager
+def cuda_stream_context(stream: torch.cuda.Stream) -> Generator[None, None, None]:
+    """Run the with-block on ``stream``, ordered after the device's current stream.
+
+    Caller must afterward sync ``stream`` back to the default stream (e.g.
+    ``torch.cuda.current_stream(stream.device).wait_stream(stream)`` from the
+    main thread once worker threads have joined) so subsequent default-stream
+    work observes this stream's results.
+
+    Args:
+        stream (torch.cuda.Stream): The pre-created per-group stream.
+
+    Example:
+        with cuda_stream_context(state_lists[CUDA_STREAM]):
+            run_step_body()
+
+    """
+    stream.wait_stream(torch.cuda.current_stream(stream.device))
+    with torch.cuda.stream(stream):
+        yield
+
+
 @cache
 def get_device_mesh(
     device_type: str,
@@ -52,3 +78,42 @@ def get_device_mesh(
 
     """
     return DeviceMesh(device_type=device_type, mesh=mesh, mesh_dim_names=mesh_dim_names)
+
+
+def create_hybrid_shard_process_groups(
+    device_mesh: DeviceMesh,
+    dist_group_size: int,
+) -> dist.ProcessGroup:
+    """Create comms process group from a hybrid shard device mesh.
+
+    Splits replicated rank groups into sub-groups of size dist_group_size
+    and returns the process group for the current rank's sub-group along
+    the shard dimension.
+
+    Args:
+        device_mesh: The hybrid shard device mesh with (replicate, shard) dims.
+        dist_group_size: Size of each distribution sub-group.
+
+    Returns:
+        The process group for communication along the shard dimension.
+    """
+    mesh_dim_names = device_mesh.mesh_dim_names
+    assert mesh_dim_names is not None, "DeviceMesh must have mesh_dim_names"
+    shard_dim_name = mesh_dim_names[1]
+    ranks_in_all_replicated_groups = device_mesh.mesh.T
+    comms_dist_group: dist.ProcessGroup | None = None
+    for ranks_in_replicated_group in ranks_in_all_replicated_groups:
+        sub_mesh = get_device_mesh(
+            device_type=device_mesh.device_type,
+            mesh=tuple(
+                map(
+                    partial(tuple),
+                    ranks_in_replicated_group.view(-1, dist_group_size).tolist(),
+                )
+            ),
+            mesh_dim_names=mesh_dim_names,
+        )
+        if dist.get_rank() in ranks_in_replicated_group:
+            comms_dist_group = sub_mesh.get_group(shard_dim_name)
+    assert comms_dist_group is not None
+    return comms_dist_group
diff --git a/distributed_shampoo/distributor/shampoo_distributor.py b/distributed_shampoo/distributor/shampoo_distributor.py
index 744fdc3..645901d 100644
--- a/distributed_shampoo/distributor/shampoo_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_distributor.py
@@ -223,12 +223,11 @@ def _merge_and_block_parameters(self) -> None:
 
         NOTE: FSDP may modify this function.
         """
-        self._global_blocked_params: tuple[Tensor, ...]
-        self._global_num_blocks_per_param: tuple[int, ...]
-        self._global_blocked_params, self._global_num_blocks_per_param = map(  # type: ignore[assignment]
-            partial(tuple),
-            self._merge_and_block_with_params(params=self._get_params_or_grads()),
+        blocked_params, num_blocks_per_param = self._merge_and_block_with_params(
+            params=self._get_params_or_grads()
         )
+        self._global_blocked_params: tuple[Tensor, ...] = tuple(blocked_params)
+        self._global_num_blocks_per_param: tuple[int, ...] = tuple(num_blocks_per_param)
 
     @abstractmethod
     def merge_and_block_gradients(
diff --git a/distributed_shampoo/distributor/shampoo_fsdp_utils.py b/distributed_shampoo/distributor/shampoo_fsdp_utils.py
index 6b18d18..7855a44 100644
--- a/distributed_shampoo/distributor/shampoo_fsdp_utils.py
+++ b/distributed_shampoo/distributor/shampoo_fsdp_utils.py
@@ -135,6 +135,7 @@ def partition_param_list(
             for cur_dict in original_params:
                 cur_fsdp_params, cur_hsdp_params, cur_other_params = (
                     partition_param_list(
+                        # pyrefly: ignore [bad-index]
                         original_params=cur_dict[
                             "params"  # type: ignore
                         ]
@@ -147,6 +148,7 @@ def partition_param_list(
 
         # Case 3: The original params is a Iterable[tuple[str, torch.Tensor]]
         case tuple():
+            # pyrefly: ignore [no-matching-overload]
             original_params_dict: dict[str, torch.Tensor] = dict(
                 original_params  # type: ignore
             )
diff --git a/distributed_shampoo/distributor/shampoo_hsdp_distributor.py b/distributed_shampoo/distributor/shampoo_hsdp_distributor.py
index a1651ca..7c617bd 100644
--- a/distributed_shampoo/distributor/shampoo_hsdp_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_hsdp_distributor.py
@@ -195,6 +195,7 @@ def __init__(
             load_balancing_config=LoadBalancingConfig(),
         )
 
+        # pyrefly: ignore [bad-override-mutable-attribute]
         self._local_block_info_list: tuple[DTensorBlockInfo, ...] = (
             self._construct_local_block_info_list(
                 group_source_ranks=tuple(
@@ -865,7 +866,7 @@ def _allocate_zeros_distributed_tensor(
         device_mesh_2d = get_device_mesh(
             device_type=device.type,
             # NOTE: Use itertools.batched(ranks_in_replicated_group, self._dist_group_size) when downstream applications are Python 3.12+ available
-            mesh=tuple(batched(ranks_in_replicated_group, self._dist_group_size)),
+            mesh=tuple(batched(ranks_in_replicated_group, self._dist_group_size)),  # noqa: B911
             mesh_dim_names=("replicate", "shard"),
         )
         # NOTE: We get all submeshes along the "replicate" dimension, then pick out
diff --git a/distributed_shampoo/distributor/shampoo_hybrid_shard_distributor.py b/distributed_shampoo/distributor/shampoo_hybrid_shard_distributor.py
index 8617efd..2f1425d 100644
--- a/distributed_shampoo/distributor/shampoo_hybrid_shard_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_hybrid_shard_distributor.py
@@ -16,6 +16,7 @@
 import torch
 from distributed_shampoo.distributor.shampoo_block_info import DTensorBlockInfo
 from distributed_shampoo.distributor.shampoo_dist_utils import (
+    create_hybrid_shard_process_groups,
     get_device_mesh,
     shampoo_comm_profiler,
 )
@@ -145,33 +146,11 @@ def __init__(
         # Create flag for distributing parameters instead of search directions.
         self._communicate_params: bool = distributed_config.communicate_params
 
-        # Initialize _dist_group and _group_rank.
-        # Note that this requires initializing all process groups.
-        # Splits replicated ranks group into smaller groups of size self._dist_group_size.
-        # Instantiates this by using DeviceMesh.
-        ranks_in_all_replicated_groups = self._hybrid_shard_device_mesh.mesh.T
-        for ranks_in_replicated_group in ranks_in_all_replicated_groups:
-            device_mesh = get_device_mesh(
-                device_type=self._hybrid_shard_device_mesh.device_type,
-                mesh=tuple(
-                    map(
-                        partial(tuple),
-                        ranks_in_replicated_group.view(
-                            -1, self._dist_group_size
-                        ).tolist(),
-                    )
-                ),
-                mesh_dim_names=("replicate", "shard"),
-            )
-            if dist.get_rank() in ranks_in_replicated_group:
-                # NOTE: We want the process group in the device mesh that the current rank
-                # belongs to but solely along the "shard" dimension for communications.
-                #
-                # For example, if the current rank is 11, then I want the process group
-                # that contains the ranks [3, 11, 19].
-                self._comms_dist_group: dist.ProcessGroup = device_mesh.get_group(
-                    "shard"
-                )
+        # Create comms process group for distributing computation across ranks.
+        self._comms_dist_group: dist.ProcessGroup = create_hybrid_shard_process_groups(
+            device_mesh=self._hybrid_shard_device_mesh,
+            dist_group_size=self._dist_group_size,
+        )
 
         comms_group_rank: int = dist.get_rank(self._comms_dist_group)
 
@@ -187,6 +166,7 @@ def __init__(
             load_balancing_config=LoadBalancingConfig(),
         )
 
+        # pyrefly: ignore [bad-override-mutable-attribute]
         self._local_block_info_list: tuple[DTensorBlockInfo, ...] = (
             self._construct_local_block_info_list(
                 group_source_ranks=tuple(
@@ -596,7 +576,7 @@ def _allocate_zeros_distributed_tensor(
         )
         device_mesh_2d = get_device_mesh(
             device_type=device.type,
-            mesh=tuple(batched(ranks_in_replicated_group, self._dist_group_size)),
+            mesh=tuple(batched(ranks_in_replicated_group, self._dist_group_size)),  # noqa: B911
             mesh_dim_names=("replicate", "shard"),
         )
         # NOTE: We get all submeshes along the "replicate" dimension, then pick out
@@ -605,7 +585,7 @@ def _allocate_zeros_distributed_tensor(
         # For the example above, this would give me submeshes [[3, 27], [11, 35], [19, 43]].
         # Note that the group source rank must belong to {0, 1, 2} in this case.
         # Suppose the group_source_rank = 1, then this would get the submesh [11, 35].
-        replicate_submesh = device_mesh_2d._get_all_submeshes(  # type: ignore[attr-defined]
+        replicate_submesh = device_mesh_2d._get_all_submeshes(
             mesh_dim_name="replicate"
         )[group_source_rank]
 
diff --git a/distributed_shampoo/distributor/tests/shampoo_distributor_test.py b/distributed_shampoo/distributor/tests/shampoo_distributor_test.py
index fcaf055..6cc90b8 100644
--- a/distributed_shampoo/distributor/tests/shampoo_distributor_test.py
+++ b/distributed_shampoo/distributor/tests/shampoo_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from typing import cast
 
@@ -45,6 +43,7 @@ def test_update_params(self) -> None:
         linear_layers: nn.ModuleList = cast(nn.ModuleList, self._model.linear_layers)
         layer_weight: torch.Tensor = cast(torch.Tensor, linear_layers[0].weight)
         layer_weight.grad = torch.ones_like(layer_weight)
+        # pyrefly: ignore [bad-argument-type, missing-attribute]
         linear_layers[0].bias.grad = None
         self._distributor.merge_and_block_gradients()
 
@@ -85,6 +84,7 @@ def test_local_grad_selector(self) -> None:
         linear_layers: nn.ModuleList = cast(nn.ModuleList, self._model.linear_layers)
         layer_weight: torch.Tensor = cast(torch.Tensor, linear_layers[0].weight)
         layer_weight.grad = torch.ones_like(layer_weight)
+        # pyrefly: ignore [bad-argument-type, missing-attribute]
         linear_layers[0].bias.grad = None
         self._distributor.merge_and_block_gradients()
 
@@ -148,6 +148,7 @@ def test_merge_and_block_gradients(self) -> None:
         linear_layers: nn.ModuleList = cast(nn.ModuleList, self._model.linear_layers)
         layer_weight: torch.Tensor = cast(torch.Tensor, linear_layers[0].weight)
         layer_weight.grad = torch.ones_like(layer_weight)
+        # pyrefly: ignore [bad-argument-type, missing-attribute]
         linear_layers[0].bias.grad = None
         actual_local_masked_block_grads = self._distributor.merge_and_block_gradients()
         expected_local_masked_block_grads = (
@@ -170,6 +171,7 @@ def test_enable_eager_nan_check(self) -> None:
         linear_layers: nn.ModuleList = cast(nn.ModuleList, self._model.linear_layers)
         layer_weight: torch.Tensor = cast(torch.Tensor, linear_layers[0].weight)
         layer_weight.grad = torch.ones_like(layer_weight)
+        # pyrefly: ignore [bad-argument-type, missing-attribute]
         linear_layers[0].bias.grad = None
 
         with unittest.mock.patch.object(torch, "isfinite") as mock_isfinite:
@@ -187,6 +189,7 @@ def test_disable_eager_nan_check(self) -> None:
         linear_layers: nn.ModuleList = cast(nn.ModuleList, self._model.linear_layers)
         layer_weight: torch.Tensor = cast(torch.Tensor, linear_layers[0].weight)
         layer_weight.grad = torch.ones_like(layer_weight)
+        # pyrefly: ignore [bad-argument-type, missing-attribute]
         linear_layers[0].bias.grad = None
 
         with unittest.mock.patch.object(torch, "isfinite") as mock_isfinite:
diff --git a/distributed_shampoo/distributor/tests/shampoo_fsdp_distributor_test.py b/distributed_shampoo/distributor/tests/shampoo_fsdp_distributor_test.py
index 18f0746..abeba64 100644
--- a/distributed_shampoo/distributor/tests/shampoo_fsdp_distributor_test.py
+++ b/distributed_shampoo/distributor/tests/shampoo_fsdp_distributor_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import logging
 import re
 import unittest
diff --git a/distributed_shampoo/examples/__init__.py b/distributed_shampoo/examples/__init__.py
index e69de29..9071dd2 100644
--- a/distributed_shampoo/examples/__init__.py
+++ b/distributed_shampoo/examples/__init__.py
@@ -0,0 +1,8 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
diff --git a/distributed_shampoo/examples/cifar10_example.py b/distributed_shampoo/examples/cifar10_example.py
index 6980850..253743f 100644
--- a/distributed_shampoo/examples/cifar10_example.py
+++ b/distributed_shampoo/examples/cifar10_example.py
@@ -1,6 +1,14 @@
 #!/usr/bin/env python3
 
-"""CIFAR-10 training example.
+
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+CIFAR-10 training example.
 
 Supports single GPU and distributed training with various parallelism strategies.
 Optimizers (sgd, adam, adamw, shampoo) can be combined with any parallelism strategy.
@@ -89,12 +97,16 @@ def main(cfg: DictConfig) -> None:
         device_mesh = create_device_mesh(cfg.dp_replicate_degree, world_size)
 
     wrapped_model = parallelism.wrap_model(model, local_rank, cfg.backend, device_mesh)
+    # FSDPModule is a runtime mixin on nn.Module, so this assert always holds;
+    # it narrows the Union for the downstream APIs that expect plain nn.Module.
+    wrapped_module = wrapped_model.model
+    assert isinstance(wrapped_module, torch.nn.Module)
 
     optimizer = instantiate_optimizer(
-        cfg, wrapped_model.model.parameters(), wrapped_model.distributed_config
+        cfg, wrapped_module.parameters(), wrapped_model.distributed_config
     )
 
-    load_checkpoint(cfg.checkpoint_dir, wrapped_model.model, optimizer)
+    load_checkpoint(cfg.checkpoint_dir, wrapped_module, optimizer)
 
     batch_size = (
         cfg.local_batch_size if parallelism.requires_distributed else cfg.batch_size
@@ -104,7 +116,7 @@ def main(cfg: DictConfig) -> None:
     )
 
     train_model(
-        model=wrapped_model.model,
+        model=wrapped_module,
         world_size=world_size,
         loss_fn=loss_fn,
         sampler=sampler if parallelism.requires_distributed else None,
diff --git a/distributed_shampoo/examples/loss_metrics.py b/distributed_shampoo/examples/loss_metrics.py
index 207a3b8..9a87264 100644
--- a/distributed_shampoo/examples/loss_metrics.py
+++ b/distributed_shampoo/examples/loss_metrics.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 Copyright (c) Meta Platforms, Inc. and affiliates.
 All rights reserved.
@@ -7,8 +9,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import logging
 
 import torch
diff --git a/distributed_shampoo/examples/parallelism.py b/distributed_shampoo/examples/parallelism.py
index a8c5678..b5a76c2 100644
--- a/distributed_shampoo/examples/parallelism.py
+++ b/distributed_shampoo/examples/parallelism.py
@@ -17,9 +17,13 @@
     compile_fsdp_parameter_metadata,
 )
 from torch import nn
-from torch.distributed._composable.fsdp import fully_shard
 from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
+from torch.distributed.fsdp import (
+    FSDPModule,
+    fully_shard,
+    FullyShardedDataParallel as FSDP,
+    ShardingStrategy,
+)
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 
@@ -27,7 +31,7 @@
 class WrappedModel:
     """Result of wrapping a model for distributed training."""
 
-    model: nn.Module
+    model: nn.Module | FSDPModule
     distributed_config: DistributedConfig | None = None
 
 
@@ -174,7 +178,7 @@ def wrap_model(
         device_mesh: DeviceMesh | None = None,
     ) -> WrappedModel:
         config = self.distributed_config() if self.distributed_config else None
-        return WrappedModel(model=fully_shard(model), distributed_config=config)  # type: ignore[arg-type]
+        return WrappedModel(model=fully_shard(model), distributed_config=config)
 
 
 @dataclass
@@ -203,6 +207,6 @@ def wrap_model(
         if self.distributed_config:
             config = self.distributed_config(device_mesh=device_mesh)
         return WrappedModel(
-            model=fully_shard(model, mesh=device_mesh),  # type: ignore[arg-type]
+            model=fully_shard(model, mesh=device_mesh),
             distributed_config=config,
         )
diff --git a/distributed_shampoo/examples/tests/convnet_test.py b/distributed_shampoo/examples/tests/convnet_test.py
index 986937a..5eb4442 100644
--- a/distributed_shampoo/examples/tests/convnet_test.py
+++ b/distributed_shampoo/examples/tests/convnet_test.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 Copyright (c) Meta Platforms, Inc. and affiliates.
 All rights reserved.
@@ -7,8 +9,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 
 import torch
diff --git a/distributed_shampoo/examples/tests/parallelism_test.py b/distributed_shampoo/examples/tests/parallelism_test.py
index d421412..0f7543f 100644
--- a/distributed_shampoo/examples/tests/parallelism_test.py
+++ b/distributed_shampoo/examples/tests/parallelism_test.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 Copyright (c) Meta Platforms, Inc. and affiliates.
 All rights reserved.
@@ -7,8 +9,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from unittest.mock import MagicMock, patch
 
diff --git a/distributed_shampoo/examples/tests/resolvers_test.py b/distributed_shampoo/examples/tests/resolvers_test.py
index bc8c359..debb79a 100644
--- a/distributed_shampoo/examples/tests/resolvers_test.py
+++ b/distributed_shampoo/examples/tests/resolvers_test.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 Copyright (c) Meta Platforms, Inc. and affiliates.
 All rights reserved.
@@ -7,8 +9,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 
 import torch
diff --git a/distributed_shampoo/examples/tests/utils_test.py b/distributed_shampoo/examples/tests/utils_test.py
index b73ef6e..54bfe58 100644
--- a/distributed_shampoo/examples/tests/utils_test.py
+++ b/distributed_shampoo/examples/tests/utils_test.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 Copyright (c) Meta Platforms, Inc. and affiliates.
 All rights reserved.
@@ -7,8 +9,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import logging
 import os
 import random
@@ -169,7 +169,8 @@ class LoadCheckpointTest(unittest.TestCase):
     def test_load_checkpoint_returns_early_if_no_checkpoint_dir(self) -> None:
         """Test that load_checkpoint returns early if checkpoint_dir is None."""
         model = nn.Linear(10, 5)
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+        # Citrine C2: foreach=True for multi-tensor execution
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.01, foreach=True)
 
         # Should not raise - just returns early
         load_checkpoint(None, model, optimizer)
@@ -177,7 +178,8 @@ def test_load_checkpoint_returns_early_if_no_checkpoint_dir(self) -> None:
     def test_load_checkpoint_returns_early_if_not_shampoo(self) -> None:
         """Test that load_checkpoint returns early for non-Shampoo optimizers."""
         model = nn.Linear(10, 5)
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+        # Citrine C2: foreach=True for multi-tensor execution
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.01, foreach=True)
 
         # Should not raise - just returns early
         load_checkpoint("/some/path", model, optimizer)
diff --git a/distributed_shampoo/examples/utils.py b/distributed_shampoo/examples/utils.py
index cc789e4..3c25cff 100644
--- a/distributed_shampoo/examples/utils.py
+++ b/distributed_shampoo/examples/utils.py
@@ -73,6 +73,7 @@ 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)
@@ -149,7 +150,9 @@ def get_data_loader_and_sampler(
     data_path = Path(data_path) / str(rank)
 
     with importlib.resources.path(
-        __package__, CIFAR_10_DATASET_FILENAME
+        # pyrefly: ignore [bad-argument-type]
+        __package__,
+        CIFAR_10_DATASET_FILENAME,
     ) as resource_path:
         if resource_path.exists():
             data_path.mkdir(parents=True, exist_ok=True)
@@ -158,14 +161,19 @@ def get_data_loader_and_sampler(
     dataset = datasets.CIFAR10(
         data_path, train=True, download=True, transform=transform
     )
-    sampler: torch.utils.data.distributed.DistributedSampler = (
-        torch.utils.data.distributed.DistributedSampler(
-            dataset, num_replicas=world_size, rank=rank, shuffle=True
-        )
+    sampler: torch.utils.data.distributed.DistributedSampler[
+        torch.utils.data.Dataset
+    ] = torch.utils.data.distributed.DistributedSampler(
+        dataset, num_replicas=world_size, rank=rank, shuffle=True
     )
     return (
         torch.utils.data.DataLoader(
-            dataset, batch_size=batch_size, sampler=sampler, num_workers=2
+            dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            num_workers=2,
+            # Citrine C0: pin_memory=True for efficient CPU-to-GPU transfer
+            pin_memory=True,
         ),
         sampler,
     )
diff --git a/distributed_shampoo/gpu_tests/iterate_averaging_test.py b/distributed_shampoo/gpu_tests/iterate_averaging_test.py
index 277a24c..cf3c932 100644
--- a/distributed_shampoo/gpu_tests/iterate_averaging_test.py
+++ b/distributed_shampoo/gpu_tests/iterate_averaging_test.py
@@ -1,6 +1,11 @@
-#!/usr/bin/env python3
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
 
-"""Tests for validating iterate averaging (GPA and Schedule-Free) equivalence.
+Tests for validating iterate averaging (GPA and Schedule-Free) equivalence.
 
 This module contains tests that validate the theoretical equivalence between Shampoo's
 iterate averaging implementations and:
@@ -27,6 +32,7 @@
 from distributed_shampoo.distributed_shampoo import DistributedShampoo
 from distributed_shampoo.shampoo_types import (
     AdamPreconditionerConfig,
+    ClassicMomentumConfig,
     GeneralizedPrimalAveragingConfig,
     IterateAveragingConfig,
     ScheduleFreeConfig,
@@ -230,3 +236,67 @@ def test_iterate_averaging_vs_gpa_adamw(
             model_linear_layers_dims=(10, 10),
             device=device,
         )
+
+    @staticmethod
+    def _shampoo_with_classic_momentum_factory(
+        iterate_averaging_config: ClassicMomentumConfig,
+        lr: float = 0.01,
+        weight_decay: float = 0.0,
+    ) -> partial[torch.optim.Optimizer]:
+        """Create a Shampoo optimizer factory with ClassicMomentumConfig and SGD preconditioner."""
+        return partial(
+            IterateAveragingTest._optim_factory,
+            optim_cls=DistributedShampoo,
+            lr=lr,
+            weight_decay=weight_decay,
+            betas=(0.0, 1.0),
+            epsilon=1e-10,
+            max_preconditioner_dim=10,
+            precondition_frequency=1,
+            start_preconditioning_step=-1,
+            weight_decay_type=WeightDecayType.L2,
+            preconditioner_config=SGDPreconditionerConfig(),  # type: ignore[abstract]
+            grafting_config=None,
+            iterate_averaging_config=iterate_averaging_config,
+        )
+
+    @parametrize("device", available_devices)
+    @parametrize("use_nesterov", [False, True])
+    def test_classic_momentum_vs_sgd_momentum(
+        self,
+        device: torch.device,
+        use_nesterov: bool,
+    ) -> None:
+        """Test that Shampoo with ClassicMomentumConfig produces equivalent results to SGD with momentum.
+
+        Uses dampening=0 so that the first-step behavior matches PyTorch SGD exactly
+        (both initialize the momentum buffer to the raw gradient when dampening=0).
+        """
+        lr = 0.01
+        momentum = 0.9
+
+        control_optim_factory = partial(
+            IterateAveragingTest._optim_factory,
+            optim_cls=SGD,
+            lr=lr,
+            weight_decay=0.0,
+            momentum=momentum,
+            dampening=0.0,
+            nesterov=use_nesterov,
+        )
+
+        experimental_optim_factory = self._shampoo_with_classic_momentum_factory(
+            iterate_averaging_config=ClassicMomentumConfig(
+                momentum=momentum,
+                dampening=0.0,
+                use_nesterov=use_nesterov,
+            ),
+            lr=lr,
+        )
+
+        compare_two_optimizers_on_weight_and_loss(
+            control_optim_factory=control_optim_factory,
+            experimental_optim_factory=experimental_optim_factory,
+            model_linear_layers_dims=(10, 10),
+            device=device,
+        )
diff --git a/distributed_shampoo/gpu_tests/shampoo_eigenvalue_correction_test.py b/distributed_shampoo/gpu_tests/shampoo_eigenvalue_correction_test.py
index a08ca9f..ebe8202 100644
--- a/distributed_shampoo/gpu_tests/shampoo_eigenvalue_correction_test.py
+++ b/distributed_shampoo/gpu_tests/shampoo_eigenvalue_correction_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import math
 import unittest
 from functools import partial
diff --git a/distributed_shampoo/gpu_tests/shampoo_grafting_test.py b/distributed_shampoo/gpu_tests/shampoo_grafting_test.py
index d4a44e5..0313ba4 100644
--- a/distributed_shampoo/gpu_tests/shampoo_grafting_test.py
+++ b/distributed_shampoo/gpu_tests/shampoo_grafting_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import math
 import unittest
 from functools import partial
diff --git a/distributed_shampoo/gpu_tests/shampoo_pt2_test.py b/distributed_shampoo/gpu_tests/shampoo_pt2_test.py
index 01822f0..8b79b08 100644
--- a/distributed_shampoo/gpu_tests/shampoo_pt2_test.py
+++ b/distributed_shampoo/gpu_tests/shampoo_pt2_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Callable
 from functools import partial
diff --git a/distributed_shampoo/gpu_tests/weight_decay_test.py b/distributed_shampoo/gpu_tests/weight_decay_test.py
index 0d20bc6..87bf224 100644
--- a/distributed_shampoo/gpu_tests/weight_decay_test.py
+++ b/distributed_shampoo/gpu_tests/weight_decay_test.py
@@ -1,6 +1,10 @@
-#!/usr/bin/env python3
-
 """
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
 Tests for weight decay types in Distributed Shampoo.
 
 This test module focuses exclusively on testing the different weight decay strategies:
diff --git a/distributed_shampoo/preconditioner/__init__.py b/distributed_shampoo/preconditioner/__init__.py
index e69de29..9071dd2 100644
--- a/distributed_shampoo/preconditioner/__init__.py
+++ b/distributed_shampoo/preconditioner/__init__.py
@@ -0,0 +1,8 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
diff --git a/distributed_shampoo/preconditioner/matrix_functions.py b/distributed_shampoo/preconditioner/matrix_functions.py
index 65ad1c8..48347e4 100644
--- a/distributed_shampoo/preconditioner/matrix_functions.py
+++ b/distributed_shampoo/preconditioner/matrix_functions.py
@@ -144,28 +144,35 @@ def wrapper(A: Tensor, *args: object, **kwargs: object) -> _FuncReturnType:
 
 def _matrix_perturbation(
     A: Tensor,
-    epsilon: float = 0.0,
+    epsilon: float | Tensor = 0.0,
     is_eigenvalues: bool = True,
 ) -> Tensor:
     """Add epsilon * I to matrix (if square) or epsilon (if vector).
 
     Args:
         A (Tensor): Matrix of interest.
-        epsilon (float): Value to add to matrix for perturbation/regularization. (Default: 0.0)
+        epsilon (float | Tensor): Value to add to matrix for perturbation/regularization.
+            May be a Python float or a 0-d Tensor; a Tensor epsilon avoids the host-device
+            sync that ``float(epsilon)`` would force on the higher-order Newton path. (Default: 0.0)
         is_eigenvalues (bool): Whether A is a matrix of eigenvalues (true) or a full matrix (false). In the former case (true), add epsilon to all values; in the latter (false), add epsilon along the diagonal. (Default: True)
 
     Returns:
         A_ridge (Tensor): Matrix with perturbation/regularization.
 
     """
+    # Fast path only when epsilon is the literal Python 0; comparing a Tensor
+    # epsilon against 0 would force a host-device sync.
+    if isinstance(epsilon, (int, float)) and epsilon == 0:
+        return A
+    if is_eigenvalues:
+        return A + epsilon
+    identity = torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
+    # `A.add(I, alpha=epsilon)` requires a Python scalar; use the elementwise
+    # form so a Tensor epsilon works without a sync.
     return (
-        (
-            A.add(torch.eye(A.shape[0], dtype=A.dtype, device=A.device), alpha=epsilon)
-            if not is_eigenvalues
-            else A + epsilon
-        )
-        if epsilon != 0
-        else A  # Fast path when epsilon is 0.0, return A without modification
+        A + epsilon * identity
+        if isinstance(epsilon, Tensor)
+        else A.add(identity, alpha=epsilon)
     )
 
 
@@ -201,7 +208,7 @@ def compute_eigenvalue_threshold(
         L: Tensor,
         rank_rtol: float | None = None,
         rank_atol: float = 0.0,
-    ) -> float:
+    ) -> Tensor:
         """Compute threshold for filtering eigenvalues based on numerical rank.
 
         Determines which eigenvalues should be considered numerically zero based on
@@ -214,13 +221,16 @@ def compute_eigenvalue_threshold(
             rank_atol (float): Absolute tolerance for determining numerical rank. (Default: 0.0)
 
         Returns:
-            threshold (float): Threshold value below which eigenvalues are treated as zero.
+            threshold (Tensor): 0-d tensor threshold below which eigenvalues are treated as zero.
         """
         if rank_rtol is None:
             rtol = L.numel() * torch.finfo(L.dtype).eps
         else:
             rtol = rank_rtol
-        return max(rank_atol, rtol * L.max().relu().item())
+        # Keep as 0-d tensor — `.item()` would force a host-device sync per
+        # eigendecomposition. `torch.where(L <= spectrum_cutoff, ...)` below
+        # accepts a 0-d tensor as the threshold.
+        return torch.clamp(rtol * L.max().relu(), min=rank_atol)
 
     match rank_deficient_stability_config:
         case PseudoInverseConfig():
@@ -270,7 +280,7 @@ def compute_eigenvalue_threshold(
 
 
 @_check_square_matrix
-def matrix_inverse_root(
+def matrix_inverse_root(  # noqa: C901
     A: Tensor,
     root: Fraction,
     root_inv_config: RootInvConfig = DefaultEigenConfig,
@@ -501,14 +511,14 @@ def matrix_inverse_root_higher_order(
                 )
 
             # develop the b coefficients array first (ref: Lakic's paper)
-            b = torch.zeros(order, dtype=A.dtype, device=A.device)
-            b[0] = 1
+            # Keep b as a Python list of floats in order to prevent host-device sync.
+            b: list[float] = [1.0]
             num = 1
             denom = 1
             for i in range(1, order):
                 num *= 1 + (i - 1) * p
                 denom *= i * p
-                b[i] = num / denom
+                b.append(num / denom)
 
             # initialize iteration, dimension, and s
             iteration = 0
@@ -528,23 +538,27 @@ def matrix_inverse_root_higher_order(
                     "Input matrix has entries close to inf, exiting root inverse"
                 )
 
-            # Now scale and setup our variables
-            epsilon = max(rel_epsilon * lambda_max_approx, abs_epsilon)
+            # Now scale and setup our variables.
+            epsilon = torch.clamp(lambda_max_approx * rel_epsilon, min=abs_epsilon)
             identity = torch.eye(n, dtype=dtype, device=A.device)
             A_ridge = _matrix_perturbation(A, epsilon=epsilon, is_eigenvalues=False)
-            lambda_max_approx += epsilon
+            lambda_max_approx = lambda_max_approx + epsilon
 
             # Figure out a constant that gives good starting location
             # We stick to a conservative setting that gives very good accuracy
             # For a ref, see https://github.com/google-research/google-research/blob/master/scalable_shampoo/pytorch/matrix_functions.py#L114
-            z = 1.0 / torch.trace(A_ridge).item()
-            X = (z ** (-s)) * identity
+            # Keep z as a 0-d tensor — `.item()` would force a host-device sync per
+            # higher-order Newton call. Use `torch.pow` / `torch.reciprocal` for
+            # tensor-native math (pyre rejects `Tensor ** float` and `float / Tensor`).
+            z = torch.reciprocal(torch.trace(A_ridge))
+            X = torch.pow(z, -s) * identity
             M = z * A_ridge
             error = torch.linalg.vector_norm(M - identity, torch.inf)
             t_iter_end = time.perf_counter()
-            logger.debug(
-                f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
-            )
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
+                )
 
             # Do one iteration of basic Newton first. This is used to mathematically guarantee convergence of higher order method.
             # TODO: we may be able to get rid of this with a more careful analysis of the convergence region
@@ -556,9 +570,10 @@ def matrix_inverse_root_higher_order(
             n_matmul = math.ceil(math.log2(p)) + 2
             iteration += 1
             t_iter_end = time.perf_counter()
-            logger.debug(
-                f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
-            )
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
+                )
 
             # main while loop
             while error > tolerance and iteration < max_iterations:
@@ -567,11 +582,9 @@ def matrix_inverse_root_higher_order(
 
                 # create M_p via Horner's rule
                 base_matrix = identity - M
-                M_p = base_matrix.mul(b[order - 1]).add_(
-                    identity, alpha=float(b[order - 2])
-                )
+                M_p = base_matrix.mul(b[order - 1]).add_(identity, alpha=b[order - 2])
                 for i in reversed(range(order - 2)):
-                    M_p = torch.addmm(identity, M_p, base_matrix, beta=float(b[i]))
+                    M_p = torch.addmm(identity, M_p, base_matrix, beta=b[i])
 
                 # rest is same as Newton
                 X = X @ M_p
@@ -581,18 +594,20 @@ def matrix_inverse_root_higher_order(
 
                 # TODO: 1.2 is the value from the Google code, can be tuned
                 if new_error > error * 1.2 or (new_error == error and error < 1e-3):
-                    logger.debug(
-                        f"Coupled inverse Newton is stagnating or diverging based on comparing current error {new_error.item()} against last iteration's error {error.item()}."
-                        f"(We assume divergence if the new error > 1.2 * previous error, and assume stagnation if they are equal.)"
-                    )
+                    if logger.isEnabledFor(logging.DEBUG):
+                        logger.debug(
+                            f"Coupled inverse Newton is stagnating or diverging based on comparing current error {new_error.item()} against last iteration's error {error.item()}."
+                            f"(We assume divergence if the new error > 1.2 * previous error, and assume stagnation if they are equal.)"
+                        )
                     termination_flag = NewtonConvergenceFlag.EARLY_STOP
                     break
                 error = new_error
 
                 t_iter_end = time.perf_counter()
-                logger.debug(
-                    f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
-                )
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        f"Iteration dur (s): {t_iter_end - t_iter_begin}, Error (|M-I|) at iteration {iteration}: {error.item()}"
+                    )
             else:
                 # determine convergence flag based on error and tolerance because the main while loop exited with False condition.
                 termination_flag = (
@@ -618,11 +633,15 @@ def matrix_inverse_root_higher_order(
                 X = torch.linalg.matrix_power(X, q)
                 n_matmul += math.ceil(math.log2(q))
 
-            logger.debug(f"Upper bound on maximum eigenvalue: {lambda_max_approx}")
-            logger.debug(f"Number of matmuls: {n_matmul}")
-            logger.debug(f"Number of iterations: {iteration}")
-            logger.debug(f"Error before powering: {true_error}")
-            logger.debug(f"Termination Flag: {termination_flag}")
+            # `lambda_max_approx` and `true_error` are 0-d GPU tensors;
+            # interpolating them in an f-string forces a host-device sync via
+            # `Tensor.__format__`. Guard the whole block behind a level check.
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(f"Upper bound on maximum eigenvalue: {lambda_max_approx}")
+                logger.debug(f"Number of matmuls: {n_matmul}")
+                logger.debug(f"Number of iterations: {iteration}")
+                logger.debug(f"Error before powering: {true_error}")
+                logger.debug(f"Termination Flag: {termination_flag}")
 
             # If we have inf/nan in our answer also raise an arithmetic exception.
             # Usually, this is due to the powering to q > 1 which can blow up entries.
@@ -965,7 +984,7 @@ def newton_schulz(
 
         """
         # Normalize the matrix A in order to ensure spectral norm <= 1.
-        X = A / max(torch.linalg.matrix_norm(A), 1e-8)
+        X = A / torch.clamp(torch.linalg.matrix_norm(A), min=1e-8)
 
         a, b, c = coefficients
 
diff --git a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
index bba4509..d4a6732 100644
--- a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
+++ b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
@@ -8,6 +8,7 @@
 """
 
 import logging
+import math
 from abc import abstractmethod
 from collections.abc import Callable, Hashable, Mapping
 from dataclasses import asdict, dataclass, field, fields
@@ -43,7 +44,12 @@
 )
 from distributed_shampoo.utils.dict_zip_iterator import DictZipIterator
 from distributed_shampoo.utils.optimizer_modules import OptimizerModule
-from distributed_shampoo.utils.shampoo_utils import compress_list, get_dtype_size
+from distributed_shampoo.utils.shampoo_utils import (
+    compress_list,
+    get_dtype_size,
+    pack_upper_triangular,
+    unpack_upper_triangular,
+)
 from torch import Tensor
 
 
@@ -58,17 +64,60 @@
 _StateValueType: TypeAlias = dict[Hashable, _SubStateValueType]
 
 
+def _allocate_packed_eye(
+    block_info: BlockInfo,
+    dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> Tensor:
+    """Allocates a packed upper triangular identity matrix.
+
+    Creates a packed-size tensor (possibly DTensor) filled with zeros, then sets
+    the diagonal entries to 1.0 on the local tensor. This avoids fancy indexing
+    on DTensor which is not supported.
+
+    Args:
+        block_info (BlockInfo): Block info providing tensor allocation functions.
+        dim (int): The dimension of the identity matrix.
+        dtype (torch.dtype): Data type for the tensor.
+        device (torch.device): Device for the tensor.
+
+    Returns:
+        packed (Tensor): A 1D tensor (or DTensor) of size dim*(dim+1)/2
+            containing the packed upper triangular identity matrix.
+    """
+    packed = block_info.allocate_zeros_tensor(
+        size=(dim * (dim + 1) // 2,),
+        dtype=dtype,
+        device=device,
+    )
+    # In row-major upper-triangular packing, the i-th diagonal entry (i, i) sits
+    # at offset i*dim - i*(i-1)//2 (skip the first i rows which contribute
+    # d, d-1, ..., d-i+1 entries).
+    local = block_info.get_tensor(packed)
+    diagonal_indices = torch.arange(dim, device=local.device)
+    diagonal_positions = (
+        diagonal_indices * dim - diagonal_indices * (diagonal_indices - 1) // 2
+    )
+    local[diagonal_positions] = 1.0
+    return packed
+
+
 @dataclass
 class BaseShampooKroneckerFactorsState(OptimizerModule):
     """Base class for Shampoo Kronecker factors (wrapped).
 
     Attributes:
         factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the factor matrices.
+            When use_symmetric_packing is True, these are 1D tensors in packed upper triangular
+            format with size d*(d+1)/2. Otherwise, they are 2D tensors of shape (d, d).
         factor_matrix_indices (tuple[str, ...]): A tuple of strings representing the indices of the factor matrices.
+        use_symmetric_packing (bool): Whether symmetric matrices are stored in packed format.
     """
 
     factor_matrices: tuple[Tensor, ...]
     factor_matrix_indices: tuple[str, ...]
+    use_symmetric_packing: bool
 
     @classmethod
     def from_block(cls, **kwargs: Any) -> "BaseShampooKroneckerFactorsState":
@@ -79,6 +128,7 @@ def from_block(cls, **kwargs: Any) -> "BaseShampooKroneckerFactorsState":
             block_info (BlockInfo): Information about the block, including methods to allocate tensors.
             factor_matrix_dtype (torch.dtype): Data type for the factor matrices.
             preconditioned_dims (tuple[int, ...]): Dimensions for which the factor matrices are preconditioned.
+            use_symmetric_packing (bool): Whether to pack symmetric matrices.
 
         Returns:
             kronecker_factors_state (BaseShampooKroneckerFactorsState): An instance of BaseShampooKroneckerFactorsState with initialized factor matrices and indices.
@@ -86,11 +136,14 @@ def from_block(cls, **kwargs: Any) -> "BaseShampooKroneckerFactorsState":
         block_info: BlockInfo = kwargs["block_info"]
         factor_matrix_dtype: torch.dtype = kwargs["factor_matrix_dtype"]
         preconditioned_dims: tuple[int, ...] = kwargs["preconditioned_dims"]
+        use_symmetric_packing: bool = kwargs["use_symmetric_packing"]
 
         return cls(
             factor_matrices=tuple(
                 block_info.allocate_zeros_tensor(
-                    size=(dim, dim),
+                    size=(
+                        (dim * (dim + 1) // 2,) if use_symmetric_packing else (dim, dim)
+                    ),
                     dtype=factor_matrix_dtype,
                     device=block_info.param.device,
                 )
@@ -100,6 +153,7 @@ def from_block(cls, **kwargs: Any) -> "BaseShampooKroneckerFactorsState":
                 ".".join((*map(str, block_info.composable_block_ids), str(k)))
                 for k in range(len(preconditioned_dims))
             ),
+            use_symmetric_packing=use_symmetric_packing,
         )
 
     def __post_init__(self) -> None:
@@ -117,9 +171,11 @@ class BaseShampooKroneckerFactorsUnwrapped:
 
     Attributes:
         factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the factor matrices.
-            These are the Kronecker factors accumulated during optimization.
+            When use_symmetric_packing is True, these are 1D tensors in packed upper triangular
+            format with size d*(d+1)/2. Otherwise, they are 2D tensors of shape (d, d).
         factor_matrix_indices (tuple[str, ...]): A tuple of strings representing the indices of
             the factor matrices, used for identification and debugging.
+        use_symmetric_packing (bool): Whether symmetric matrices are stored in packed format.
         roots (tuple[float, ...]): A tuple of float values representing the inverse exponent roots
             to be applied to each factor matrix during preconditioner computation.
         amortized_computation_config (MatrixFunctionConfig): Configuration for the amortized
@@ -135,11 +191,13 @@ class BaseShampooKroneckerFactorsUnwrapped:
 
     factor_matrices: tuple[Tensor, ...]
     factor_matrix_indices: tuple[str, ...]
+    use_symmetric_packing: bool
     roots: tuple[float, ...]
     amortized_computation_config: MatrixFunctionConfig
     epsilon: float
     num_tolerated_failed_amortized_computations: int
     use_trace_scaling: bool = False
+    trace_scaling_exponent: float = 0.5
     _failed_amortized_computation_counter: int = field(init=False, default=0)
 
     def __post_init__(self) -> None:
@@ -164,11 +222,13 @@ def _get_field_dict(self) -> dict[str, Any]:
             for field in fields(self)
             if field.name
             not in (
+                "_failed_amortized_computation_counter",
                 "amortized_computation_config",
                 "epsilon",
                 "num_tolerated_failed_amortized_computations",
+                "use_symmetric_packing",
                 "use_trace_scaling",
-                "_failed_amortized_computation_counter",
+                "trace_scaling_exponent",
             )
         }
 
@@ -206,6 +266,51 @@ def _amortized_computation(
             exception (Exception | None): Any exception that occurred during computation, or None if successful.
         """
 
+    @staticmethod
+    def _packed_dim(packed: Tensor) -> int:
+        """Computes the matrix dimension from a packed upper triangular tensor's length.
+
+        For a d×d symmetric matrix, the packed length is d*(d+1)/2. This method
+        inverts that formula.
+
+        Args:
+            packed (Tensor): A 1D tensor in packed upper triangular format.
+
+        Returns:
+            dim (int): The dimension of the original square matrix.
+        """
+        return (-1 + math.isqrt(1 + 8 * packed.numel())) // 2
+
+    def _pack_matrix(self, matrix: Tensor) -> Tensor:
+        """Packs a 2D symmetric matrix to upper triangular form when packing is enabled.
+
+        When use_symmetric_packing is False, returns the input matrix unchanged.
+        """
+        return pack_upper_triangular(matrix) if self.use_symmetric_packing else matrix
+
+    def _unpack_matrix(self, matrix: Tensor) -> Tensor:
+        """Reconstructs a full 2D symmetric matrix from packed storage when packing is enabled.
+
+        When use_symmetric_packing is False, returns the input matrix unchanged.
+        Note: When packing is enabled, this allocates a new d×d tensor.
+        """
+        return (
+            unpack_upper_triangular(matrix, self._packed_dim(matrix))
+            if self.use_symmetric_packing
+            else matrix
+        )
+
+    def get_unpacked_factor_matrices(self) -> tuple[Tensor, ...]:
+        """Reconstructs full symmetric factor matrices from packed storage.
+
+        Note: This method allocates new d×d tensors each call, trading compute
+        for the memory savings from packed storage.
+
+        Returns:
+            factor_matrices (tuple[Tensor, ...]): A tuple of 2D symmetric tensors.
+        """
+        return tuple(map(self._unpack_matrix, self.factor_matrices))
+
     @profile_decorator
     def amortized_computation(self, bias_correction2: float) -> None:
         """Performs amortized computation for Shampoo preconditioners.
@@ -227,11 +332,11 @@ def amortized_computation(self, bias_correction2: float) -> None:
         """
         last_seen_exception: Exception | None = None
         for kronecker_factors_iter_dict in DictZipIterator(data=self._get_field_dict()):
-            bias_corrected_factor_matrix, factor_matrix_index = (
-                # Incorporate bias correction.
-                kronecker_factors_iter_dict["factor_matrices"] / bias_correction2,
-                kronecker_factors_iter_dict["factor_matrix_indices"],
+            # Incorporate bias correction, then unpack to full 2D matrix if packed.
+            bias_corrected_factor_matrix = self._unpack_matrix(
+                kronecker_factors_iter_dict["factor_matrices"] / bias_correction2
             )
+            factor_matrix_index = kronecker_factors_iter_dict["factor_matrix_indices"]
 
             # Apply trace scaling if enabled.
             # This normalizes the factor matrix by 1/sqrt(trace) before computing
@@ -261,8 +366,8 @@ def amortized_computation(self, bias_correction2: float) -> None:
             if self.use_trace_scaling:
                 trace = torch.trace(bias_corrected_factor_matrix)
                 safe_trace = torch.clamp(trace, min=self.epsilon)
-                bias_corrected_factor_matrix = (
-                    bias_corrected_factor_matrix * torch.rsqrt(safe_trace)
+                bias_corrected_factor_matrix = bias_corrected_factor_matrix * torch.pow(
+                    safe_trace, -self.trace_scaling_exponent
                 )
 
             # Check for nan or inf values.
@@ -342,8 +447,15 @@ def raise_preconditioner_value_error(
 class RootInvShampooKroneckerFactorsState(BaseShampooKroneckerFactorsState):
     """Shampoo Kronecker factors (wrapped) for storing in the optimizer state.
 
+    When use_symmetric_packing is True, the inverse factor matrices are stored in
+    packed upper triangular format to save ~50% memory, since they are symmetric
+    positive semi-definite.
+
     Attributes:
-        inv_factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the inverse of the factor matrices.
+        inv_factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the inverse
+            of the factor matrices. When use_symmetric_packing is True, these are 1D tensors
+            containing the packed upper triangular elements with size d*(d+1)/2. Otherwise,
+            they are 2D tensors of shape (d, d).
         factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the factor matrices.
         factor_matrix_indices (tuple[str, ...]): A tuple of strings representing the indices of the factor matrices.
     """
@@ -368,6 +480,7 @@ def from_block(cls, **kwargs: Any) -> "RootInvShampooKroneckerFactorsState":
             "preconditioner_config"
         ]
         preconditioned_dims: tuple[int, ...] = kwargs["preconditioned_dims"]
+        use_symmetric_packing: bool = preconditioner_config.use_symmetric_packing
 
         return cls(
             **asdict(
@@ -375,14 +488,26 @@ def from_block(cls, **kwargs: Any) -> "RootInvShampooKroneckerFactorsState":
                     block_info=block_info,
                     factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
+                    use_symmetric_packing=use_symmetric_packing,
                 )
             ),
             # Initialize inv_factor_matrices as identity matrices.
+            # When packing, allocate packed-size zeros (possibly DTensor), then fill identity
+            # values on the local tensor to avoid fancy indexing on DTensor.
             inv_factor_matrices=tuple(
-                block_info.allocate_eye_tensor(
-                    n=dim,
-                    dtype=preconditioner_config.inv_factor_matrix_dtype,
-                    device=block_info.param.device,
+                (
+                    _allocate_packed_eye(
+                        block_info=block_info,
+                        dim=dim,
+                        dtype=preconditioner_config.inv_factor_matrix_dtype,
+                        device=block_info.param.device,
+                    )
+                    if use_symmetric_packing
+                    else block_info.allocate_eye_tensor(
+                        n=dim,
+                        dtype=preconditioner_config.inv_factor_matrix_dtype,
+                        device=block_info.param.device,
+                    )
                 )
                 for dim in preconditioned_dims
             ),
@@ -400,9 +525,15 @@ class RootInvShampooKroneckerFactorsUnwrapped(BaseShampooKroneckerFactorsUnwrapp
     This class implements the Root Inverse variant of Shampoo, which directly computes
     the inverse root of factor matrices for preconditioning.
 
+    When use_symmetric_packing is True, the inverse factor matrices are stored in packed
+    upper triangular format to save ~50% memory. They are unpacked to full symmetric
+    matrices when needed for computation.
+
     Attributes:
         inv_factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the inverse
-            of the factor matrices. These are the preconditioners that are applied to gradients.
+            of the factor matrices. When use_symmetric_packing is True, these are 1D tensors
+            containing the packed upper triangular elements with size d*(d+1)/2. Otherwise,
+            they are 2D tensors of shape (d, d).
         factor_matrices (tuple[Tensor, ...]): A tuple of tensors representing the factor matrices.
             These are the Kronecker factors accumulated during optimization.
         factor_matrix_indices (tuple[str, ...]): A tuple of strings representing the indices of
@@ -432,6 +563,7 @@ def from_kronecker_factors_state(
         epsilon: float,
         num_tolerated_failed_amortized_computations: int,
         use_trace_scaling: bool = False,
+        trace_scaling_exponent: float = 0.5,
     ) -> "RootInvShampooKroneckerFactorsUnwrapped":
         """
         Constructs a RootInvShampooKroneckerFactorsUnwrapped object from the given Kronecker factors state.
@@ -456,6 +588,8 @@ def from_kronecker_factors_state(
                 failed amortized computations that can be tolerated before raising an error.
             use_trace_scaling (bool): Flag for whether to normalize the factor matrix by its trace's
                 sqrt before computing the inverse root. (Default: False)
+            trace_scaling_exponent (float): Exponent applied to the trace when use_trace_scaling is True;
+                the factor matrix is multiplied by trace ** (-trace_scaling_exponent). (Default: 0.5)
 
         Returns:
             kronecker_factors_unwrapped (RootInvShampooKroneckerFactorsUnwrapped): An instance of
@@ -473,13 +607,27 @@ def from_kronecker_factors_state(
                 map(unwrapped_tensor_getter, kronecker_factors_state.factor_matrices)
             ),
             factor_matrix_indices=kronecker_factors_state.factor_matrix_indices,
+            use_symmetric_packing=kronecker_factors_state.use_symmetric_packing,
             roots=roots,
             amortized_computation_config=amortized_computation_config,
             epsilon=epsilon,
             num_tolerated_failed_amortized_computations=num_tolerated_failed_amortized_computations,
             use_trace_scaling=use_trace_scaling,
+            trace_scaling_exponent=trace_scaling_exponent,
         )
 
+    def get_unpacked_inv_factor_matrices(self) -> tuple[Tensor, ...]:
+        """Reconstructs full symmetric inverse factor matrices from packed storage.
+
+        Note: This method allocates new d×d tensors each call. It is called every
+        precondition step (not just amortized computation steps), trading compute
+        for the memory savings from packed storage.
+
+        Returns:
+            inv_factor_matrices (tuple[Tensor, ...]): A tuple of 2D symmetric tensors.
+        """
+        return tuple(map(self._unpack_matrix, self.inv_factor_matrices))
+
     @torch.compiler.disable
     def _amortized_computation(
         self,
@@ -496,6 +644,9 @@ def _amortized_computation(
         configuration specified in amortized_computation_config. Error handling is included
         to gracefully recover from numerical issues.
 
+        When use_symmetric_packing is True, the computed inverse root is packed into upper
+        triangular format before being stored.
+
         Args:
             bias_corrected_factor_matrix (Tensor): The factor matrix after bias correction
                 has been applied.
@@ -503,7 +654,8 @@ def _amortized_computation(
                 inv_factor_matrices and roots values for the computation.
 
         Returns:
-            computed_quantities (dict[str, Tensor]): A dictionary with the computed inverse factor matrices.
+            computed_quantities (dict[str, Tensor]): A dictionary with the computed inverse
+                factor matrices, packed if use_symmetric_packing is enabled.
             exception (Exception | None): Any exception that occurred during computation, or None if successful.
 
         Note:
@@ -516,15 +668,14 @@ def _amortized_computation(
         )
 
         try:
-            # Compute inverse preconditioners
-            return {
-                "inv_factor_matrices": matrix_inverse_root(
-                    A=bias_corrected_factor_matrix,
-                    root=Fraction(root),
-                    root_inv_config=self.amortized_computation_config,
-                    epsilon=self.epsilon,
-                ).to(dtype=inv_factor_matrix.dtype)
-            }, None
+            # Compute inverse preconditioners, then pack to upper triangular if enabled.
+            computed = matrix_inverse_root(
+                A=bias_corrected_factor_matrix,
+                root=Fraction(root),
+                root_inv_config=self.amortized_computation_config,
+                epsilon=self.epsilon,
+            ).to(dtype=inv_factor_matrix.dtype)
+            return {"inv_factor_matrices": self._pack_matrix(computed)}, None
         except Exception as exception:
             return {"inv_factor_matrices": inv_factor_matrix}, exception
 
@@ -576,6 +727,7 @@ def from_block(cls, **kwargs: Any) -> "EigendecomposedShampooKroneckerFactorsSta
                     block_info=block_info,
                     factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
+                    use_symmetric_packing=preconditioner_config.use_symmetric_packing,
                 )
             ),
             # Initialize factor_matrices_eigenvectors as identity matrices.
@@ -652,6 +804,7 @@ def from_kronecker_factors_state(
         epsilon: float,
         num_tolerated_failed_amortized_computations: int,
         use_trace_scaling: bool = False,
+        trace_scaling_exponent: float = 0.5,
     ) -> "EigendecomposedShampooKroneckerFactorsUnwrapped":
         """
         Constructs an EigendecomposedShampooKroneckerFactorsUnwrapped object from the given Kronecker factors state.
@@ -675,6 +828,8 @@ def from_kronecker_factors_state(
                 failed amortized computations that can be tolerated before raising an error.
             use_trace_scaling (bool): Flag for whether to normalize the factor matrix by its trace's
                 sqrt before computing the eigendecomposition. (Default: False)
+            trace_scaling_exponent (float): Exponent applied to the trace when use_trace_scaling is True;
+                the factor matrix is multiplied by trace ** (-trace_scaling_exponent). (Default: 0.5)
 
         Returns:
             kronecker_factors_unwrapped (EigendecomposedShampooKroneckerFactorsUnwrapped): An instance of
@@ -700,11 +855,13 @@ def from_kronecker_factors_state(
                 map(unwrapped_tensor_getter, kronecker_factors_state.factor_matrices)
             ),
             factor_matrix_indices=kronecker_factors_state.factor_matrix_indices,
+            use_symmetric_packing=kronecker_factors_state.use_symmetric_packing,
             roots=roots,
             amortized_computation_config=amortized_computation_config,
             epsilon=epsilon,
             num_tolerated_failed_amortized_computations=num_tolerated_failed_amortized_computations,
             use_trace_scaling=use_trace_scaling,
+            trace_scaling_exponent=trace_scaling_exponent,
         )
 
     @torch.compiler.disable
@@ -825,6 +982,7 @@ def from_block(
                     block_info=block_info,
                     factor_matrix_dtype=preconditioner_config.factor_matrix_dtype,
                     preconditioned_dims=preconditioned_dims,
+                    use_symmetric_packing=preconditioner_config.use_symmetric_packing,
                 )
             ),
             # Initialize factor_matrices_eigenvectors as identity matrices.
@@ -898,6 +1056,7 @@ def from_kronecker_factors_state(
         epsilon: float,
         num_tolerated_failed_amortized_computations: int,
         use_trace_scaling: bool = False,
+        trace_scaling_exponent: float = 0.5,
     ) -> "EigenvalueCorrectedShampooKroneckerFactorsUnwrapped":
         """
         Constructs an EigenvalueCorrectedShampooKroneckerFactorsUnwrapped object from the given Kronecker factors state.
@@ -924,6 +1083,8 @@ def from_kronecker_factors_state(
                 failed amortized computations that can be tolerated before raising an error.
             use_trace_scaling (bool): Flag for whether to normalize the factor matrix by its trace's
                 sqrt before computing the eigendecomposition. (Default: False)
+            trace_scaling_exponent (float): Exponent applied to the trace when use_trace_scaling is True;
+                the factor matrix is multiplied by trace ** (-trace_scaling_exponent). (Default: 0.5)
 
         Returns:
             kronecker_factors_unwrapped (EigenvalueCorrectedShampooKroneckerFactorsUnwrapped): An instance of
@@ -946,11 +1107,13 @@ def from_kronecker_factors_state(
                 map(unwrapped_tensor_getter, kronecker_factors_state.factor_matrices)
             ),
             factor_matrix_indices=kronecker_factors_state.factor_matrix_indices,
+            use_symmetric_packing=kronecker_factors_state.use_symmetric_packing,
             roots=roots,
             amortized_computation_config=amortized_computation_config,
             epsilon=epsilon,
             num_tolerated_failed_amortized_computations=num_tolerated_failed_amortized_computations,
             use_trace_scaling=use_trace_scaling,
+            trace_scaling_exponent=trace_scaling_exponent,
         )
 
     @torch.compiler.disable
@@ -1197,6 +1360,7 @@ def _create_kronecker_factors_state(
                     epsilon=self._epsilon,
                     num_tolerated_failed_amortized_computations=self._preconditioner_config.num_tolerated_failed_amortized_computations,
                     use_trace_scaling=self._preconditioner_config.use_trace_scaling,
+                    trace_scaling_exponent=self._preconditioner_config.trace_scaling_exponent,
                 )
             )
 
@@ -1247,7 +1411,13 @@ def update_preconditioners(
         self._update_factor_matrices(masked_grad_list=masked_grad_list)
 
         # Update bias correction term based on step.
-        if self._use_bias_correction and self._beta2 < 1.0:
+        # When drop_weighting_factor_on_gsquare is True, _bias_correction2 stays at 1.0
+        # because the recurrence is no longer an EMA.
+        if (
+            self._use_bias_correction
+            and self._beta2 < 1.0
+            and not self._preconditioner_config.drop_weighting_factor_on_gsquare
+        ):
             self._bias_correction2 = torch.tensor(1.0) - self._beta2**step
 
         # In Shampoo, this is equivalent to computing the inverse factor matrix.
@@ -1352,8 +1522,11 @@ def _update_factor_matrices(self, masked_grad_list: tuple[Tensor, ...]) -> None:
             if not kronecker_factors.factor_matrices:
                 continue
 
-            outer_product_list = self._compute_outer_product_list(
-                grad, order, preconditioned_dims_selector, kronecker_factors
+            outer_product_list = tuple(
+                kronecker_factors._pack_matrix(op)
+                for op in self._compute_outer_product_list(
+                    grad, order, preconditioned_dims_selector, kronecker_factors
+                )
             )
 
             if self._beta2 != 1.0:
@@ -1388,16 +1561,18 @@ def _precondition_grad(
         preconditioner_list_iter = iter(preconditioner_list)
 
         return reduce(
-            lambda grad, should_precondition: torch.tensordot(
-                # Use the single target dtype for all operations
-                grad.to(dtype=target_dtype),
-                # Use the actual iterator for the operation
-                next(preconditioner_list_iter),
-                dims=dims,
-            )
-            if should_precondition
-            # Perform a left rotation on grad if not preconditioned.
-            else grad.permute(*range(1, grad.ndim), 0),
+            lambda grad, should_precondition: (
+                torch.tensordot(
+                    # Use the single target dtype for all operations
+                    grad.to(dtype=target_dtype),
+                    # Use the actual iterator for the operation
+                    next(preconditioner_list_iter),
+                    dims=dims,
+                )
+                if should_precondition
+                # Perform a left rotation on grad if not preconditioned.
+                else grad.permute(*range(1, grad.ndim), 0)
+            ),
             preconditioned_dims_selector,
             grad,
         ).to(dtype=grad.dtype)
@@ -1611,7 +1786,7 @@ def _compute_preconditioned_gradient(
         return self._precondition_grad(
             grad=grad,
             preconditioned_dims_selector=preconditioned_dims_selector,
-            preconditioner_list=kronecker_factors.inv_factor_matrices,
+            preconditioner_list=kronecker_factors.get_unpacked_inv_factor_matrices(),
         )
 
 
@@ -1779,6 +1954,9 @@ def _compute_outer_product_list(
         kronecker_factors: RootInvShampooKroneckerFactorsUnwrapped,
     ) -> tuple[Tensor, ...]:
         # Construct outer product list for updating Kronecker factors.
+        unpacked_inv_factor_matrices = (
+            kronecker_factors.get_unpacked_inv_factor_matrices()
+        )
         outer_product_list = []
         for idx_of_k, k in enumerate(
             compress_list(range(order), preconditioned_dims_selector)
@@ -1792,7 +1970,7 @@ def _compute_outer_product_list(
                 preconditioner_list=tuple(
                     inv_factor_matrix
                     for idx, inv_factor_matrix in enumerate(
-                        kronecker_factors.inv_factor_matrices
+                        unpacked_inv_factor_matrices
                     )
                     if idx != idx_of_k
                 ),
diff --git a/distributed_shampoo/preconditioner/tests/adagrad_preconditioner_list_test.py b/distributed_shampoo/preconditioner/tests/adagrad_preconditioner_list_test.py
index d6d6d02..a821f09 100644
--- a/distributed_shampoo/preconditioner/tests/adagrad_preconditioner_list_test.py
+++ b/distributed_shampoo/preconditioner/tests/adagrad_preconditioner_list_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 from collections.abc import Hashable
 from typing import Any
 
diff --git a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
index 5f1aa23..0ccf5b7 100644
--- a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
+++ b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import itertools
 import re
 import unittest
diff --git a/distributed_shampoo/preconditioner/tests/matrix_functions_types_test.py b/distributed_shampoo/preconditioner/tests/matrix_functions_types_test.py
index d08d6f6..933bbfe 100644
--- a/distributed_shampoo/preconditioner/tests/matrix_functions_types_test.py
+++ b/distributed_shampoo/preconditioner/tests/matrix_functions_types_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 
diff --git a/distributed_shampoo/preconditioner/tests/preconditioner_list_test_utils.py b/distributed_shampoo/preconditioner/tests/preconditioner_list_test_utils.py
index 9e9425b..e349833 100644
--- a/distributed_shampoo/preconditioner/tests/preconditioner_list_test_utils.py
+++ b/distributed_shampoo/preconditioner/tests/preconditioner_list_test_utils.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import abc
 import unittest
 from typing import Any
diff --git a/distributed_shampoo/preconditioner/tests/sgd_preconditioner_list_test.py b/distributed_shampoo/preconditioner/tests/sgd_preconditioner_list_test.py
index ef3109e..51f7939 100644
--- a/distributed_shampoo/preconditioner/tests/sgd_preconditioner_list_test.py
+++ b/distributed_shampoo/preconditioner/tests/sgd_preconditioner_list_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import torch
 from distributed_shampoo.preconditioner.sgd_preconditioner_list import (
     SGDPreconditionerList,
diff --git a/distributed_shampoo/preconditioner/tests/shampoo_preconditioner_list_test.py b/distributed_shampoo/preconditioner/tests/shampoo_preconditioner_list_test.py
index 3d3c5c1..90528c4 100644
--- a/distributed_shampoo/preconditioner/tests/shampoo_preconditioner_list_test.py
+++ b/distributed_shampoo/preconditioner/tests/shampoo_preconditioner_list_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import abc
 import math
 import re
@@ -77,11 +75,13 @@ class InverseRootProperties(AmortizedComputationProperties):
     """Dataclass for properties of matrix_inverse_root function."""
 
     amortized_computation_function_name: str = "matrix_inverse_root"
+    # pyrefly: ignore [bad-override-mutable-attribute]
     invalid_amortized_computation_return_values: tuple[Tensor, Tensor] = (
-        torch.tensor([torch.nan]),
-        torch.tensor([torch.inf]),
+        torch.tensor([[torch.nan]]),
+        torch.tensor([[torch.inf]]),
     )
-    valid_amortized_computation_return_value: Tensor = torch.tensor([1.0])
+    # pyrefly: ignore [bad-override-mutable-attribute]
+    valid_amortized_computation_return_value: Tensor = torch.tensor([[1.0]])
 
 
 @dataclass
@@ -89,12 +89,14 @@ class EigendecompositionProperties(AmortizedComputationProperties):
     """Dataclass for properties of matrix_eigendecomposition function."""
 
     amortized_computation_function_name: str = "matrix_eigendecomposition"
+    # pyrefly: ignore [bad-override-mutable-attribute]
     invalid_amortized_computation_return_values: tuple[
         tuple[Tensor, Tensor], tuple[Tensor, Tensor]
     ] = (
         (torch.tensor([torch.nan]), torch.tensor([torch.nan])),
         (torch.tensor([torch.inf]), torch.tensor([torch.inf])),
     )
+    # pyrefly: ignore [bad-override-mutable-attribute]
     valid_amortized_computation_return_value: tuple[Tensor, Tensor] = (
         torch.tensor([1.0]),
         torch.tensor([1.0]),
@@ -173,9 +175,9 @@ def test_adaptive_amortized_computation_frequency(self) -> None:
                     torch,
                     "sqrt",
                     side_effect=(
-                        -math.inf
-                        if v
-                        else math.inf  # Return -inf for True (skip) and inf for False (compute)
+                        (
+                            -math.inf if v else math.inf
+                        )  # Return -inf for True (skip) and inf for False (compute)
                         for v in CRITERION_RESULTS
                     ),
                 ) as mock_criterion,
@@ -563,9 +565,11 @@ def test_precondition_grad(self, dims: tuple[list[int], list[int]]) -> None:
 
             # Create a control preconditioner list, using identity matrices where not preconditioning.
             control_preconditioner_list = tuple(
-                preconditioner
-                if should_precondition
-                else torch.eye(preconditioner.shape[0])
+                (
+                    preconditioner
+                    if should_precondition
+                    else torch.eye(preconditioner.shape[0])
+                )
                 for preconditioner, should_precondition in zip(
                     preconditioner_list,
                     experimental_preconditioned_dims_selector,
@@ -617,6 +621,59 @@ def _expected_num_bytes(self) -> int:
         def _expected_compress_list_call_count(self) -> int:
             return 3
 
+        def test_drop_weighting_factor_on_gsquare_skips_bias_correction2(self) -> None:
+            # When drop_weighting_factor_on_gsquare=True, the recurrence is no longer
+            # an EMA, so _bias_correction2 must stay at 1.0 even with use_bias_correction=True
+            # and beta2 < 1.
+            beta2 = 0.9
+            preconditioner_list = self._instantiate_preconditioner_list(
+                beta2=beta2,
+                weighting_factor=1.0,
+                use_bias_correction=True,
+                preconditioner_config=replace(
+                    self._default_preconditioner_config,
+                    drop_weighting_factor_on_gsquare=True,
+                ),
+            )
+            assert isinstance(preconditioner_list, BaseShampooPreconditionerList)
+            preconditioner_list.update_preconditioners(
+                masked_grad_list=(
+                    torch.tensor([1.0, 0.0]),
+                    torch.eye(2) / math.sqrt(2.0),
+                    torch.tensor([[1.0, 0.0]]),
+                    torch.tensor(1.0),
+                ),
+                step=torch.tensor(5),
+                perform_amortized_computation=False,
+            )
+            self.assertEqual(preconditioner_list._bias_correction2.item(), 1.0)
+
+        def test_bias_correction2_applied_when_drop_weighting_factor_off(self) -> None:
+            # Sanity check: when drop_weighting_factor_on_gsquare=False (the default), the
+            # standard 1 - beta2**step formula is still applied.
+            beta2 = 0.9
+            step = 5
+            preconditioner_list = self._instantiate_preconditioner_list(
+                beta2=beta2,
+                weighting_factor=1 - beta2,
+                use_bias_correction=True,
+            )
+            assert isinstance(preconditioner_list, BaseShampooPreconditionerList)
+            preconditioner_list.update_preconditioners(
+                masked_grad_list=(
+                    torch.tensor([1.0, 0.0]),
+                    torch.eye(2) / math.sqrt(2.0),
+                    torch.tensor([[1.0, 0.0]]),
+                    torch.tensor(1.0),
+                ),
+                step=torch.tensor(step),
+                perform_amortized_computation=False,
+            )
+            torch.testing.assert_close(
+                preconditioner_list._bias_correction2,
+                torch.tensor(1.0 - beta2**step),
+            )
+
     class ClassicShampooPreconditionerListTest(BaseShampooPreconditionerListTest):
         @property
         @abc.abstractmethod
@@ -831,6 +888,43 @@ def test_update_preconditioners_and_precondition(self) -> None:
                 ),
             )
 
+            """
+            For the case of beta2 = 1, use_trace_scaling = True, and a non-default
+            trace_scaling_exponent = 1.0, the factor matrix is normalized by trace ** (-1.0)
+            (instead of the default sqrt, i.e. trace ** (-0.5)) before the inverse root.
+
+            With L = R = I (trace = 2): L' = R' = I * 2 ** (-1.0) = I / 2.
+
+            (1) Tensor of Size 2:   P = L'^{-1/2} G2 = 2^{1/2} * G2
+            (2) Tensor of Size 2 x 2: P = L'^{-1/4} G2 R'^{-1/4} = 2^{1/4} * G2 * 2^{1/4} = 2^{1/2} * G2
+            (3) Tensor of Size 1 x 2: L = 2 (scalar) -> L' = 1 -> L'^{-1/4} = 1; R' = I / 2 -> R'^{-1/4} = 2^{1/4}
+                P = 1 * G2 * 2^{1/4} = 2^{1/4} * G2
+            (4) Tensor of Size 0:   No preconditioner is applied. P = G2.
+            """
+            masked_expected_preconditioned_grad_list_trace_scaling_exponent = [
+                masked_grad_list2[0] * 2.0 ** (1 / 2),
+                masked_grad_list2[1] * 2.0 ** (1 / 2),
+                masked_grad_list2[2] * 2.0 ** (1 / 4),
+                masked_grad_list2[3],  # 0D tensor, no preconditioner
+            ]
+
+            self._verify_preconditioner_updates(
+                preconditioner_list=self._instantiate_preconditioner_list(
+                    beta2=1.0,
+                    weighting_factor=1.0,
+                    use_bias_correction=True,
+                    preconditioner_config=replace(
+                        self._default_preconditioner_config,
+                        use_trace_scaling=True,
+                        trace_scaling_exponent=1.0,
+                    ),
+                ),
+                masked_grad_lists=[masked_grad_list1, masked_grad_list2],
+                masked_expected_preconditioned_grad_list=tuple(
+                    masked_expected_preconditioned_grad_list_trace_scaling_exponent
+                ),
+            )
+
         def test_update_preconditioners_and_precondition_with_epsilon(self) -> None:
             """
             We provide examples where we deliberately choose a large epsilon. This is to ensure that
@@ -1101,6 +1195,9 @@ def _default_preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
             DefaultShampooConfig,
             factor_matrix_dtype=torch.float64,
             inv_factor_matrix_dtype=torch.float64,
+            # Pin to True so the standard packed path is exercised even if the
+            # config default changes in the future.
+            use_symmetric_packing=True,
         )
 
     @property
@@ -1129,6 +1226,9 @@ def _default_preconditioner_config(  # type: ignore[override]
             factor_matrix_dtype=torch.float64,
             factor_matrix_eigenvectors_dtype=torch.float64,
             factor_matrix_eigenvalues_dtype=torch.float64,
+            # Pin to True so the standard packed path is exercised even if the
+            # config default changes in the future.
+            use_symmetric_packing=True,
         )
 
     @property
@@ -1152,6 +1252,9 @@ def _default_preconditioner_config(
             factor_matrix_dtype=torch.float64,
             factor_matrix_eigenvectors_dtype=torch.float64,
             corrected_eigenvalues_dtype=torch.float64,
+            # Pin to True so the standard packed path is exercised even if the
+            # config default changes in the future.
+            use_symmetric_packing=True,
         )
 
     @property
@@ -1414,6 +1517,9 @@ def _default_preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
         return replace(
             RootInvKLShampooPreconditionerConfig(),
             factor_matrix_dtype=torch.float64,
+            # Pin to True so the standard packed path is exercised even if the
+            # config default changes in the future.
+            use_symmetric_packing=True,
         )
 
     @property
@@ -1438,6 +1544,9 @@ def _default_preconditioner_config(  # type: ignore[override]
             factor_matrix_dtype=torch.float64,
             factor_matrix_eigenvectors_dtype=torch.float64,
             factor_matrix_eigenvalues_dtype=torch.float64,
+            # Pin to True so the standard packed path is exercised even if the
+            # config default changes in the future.
+            use_symmetric_packing=True,
         )
 
     @property
@@ -1541,3 +1650,68 @@ def test_update_preconditioners_and_precondition_with_epsilon(self) -> None:
                 masked_expected_preconditioned_grad_list
             ),
         )
+
+
+# ---- Unpacked (use_symmetric_packing=False) test variants ----
+# These rerun all tests with symmetric packing disabled to verify the unpacked path.
+
+
+class RootInvShampooPreconditionerListUnpackedTest(
+    RootInvShampooPreconditionerListTest,
+):
+    @property
+    def _default_preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
+        return replace(
+            super()._default_preconditioner_config,
+            use_symmetric_packing=False,
+        )
+
+
+class EigendecomposedShampooPreconditionerListUnpackedTest(
+    EigendecomposedShampooPreconditionerListTest,
+):
+    @property
+    def _default_preconditioner_config(  # type: ignore[override]
+        self,
+    ) -> EigendecomposedShampooPreconditionerConfig:
+        return replace(
+            super()._default_preconditioner_config,
+            use_symmetric_packing=False,
+        )
+
+
+class EigenvalueCorrectedShampooPreconditionerListUnpackedTest(
+    EigenvalueCorrectedShampooPreconditionerListTest,
+):
+    @property
+    def _default_preconditioner_config(
+        self,
+    ) -> EigenvalueCorrectedShampooPreconditionerConfig:
+        return replace(
+            super()._default_preconditioner_config,
+            use_symmetric_packing=False,
+        )
+
+
+class RootInvKLShampooPreconditionerListUnpackedTest(
+    RootInvKLShampooPreconditionerListTest,
+):
+    @property
+    def _default_preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
+        return replace(
+            super()._default_preconditioner_config,
+            use_symmetric_packing=False,
+        )
+
+
+class EigendecomposedKLShampooPreconditionerListUnpackedTest(
+    EigendecomposedKLShampooPreconditionerListTest,
+):
+    @property
+    def _default_preconditioner_config(  # type: ignore[override]
+        self,
+    ) -> EigendecomposedKLShampooPreconditionerConfig:
+        return replace(
+            super()._default_preconditioner_config,
+            use_symmetric_packing=False,
+        )
diff --git a/distributed_shampoo/preconditioner/tests/sign_descent_preconditioner_list_test.py b/distributed_shampoo/preconditioner/tests/sign_descent_preconditioner_list_test.py
index c9dbae2..9958603 100644
--- a/distributed_shampoo/preconditioner/tests/sign_descent_preconditioner_list_test.py
+++ b/distributed_shampoo/preconditioner/tests/sign_descent_preconditioner_list_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 from collections.abc import Callable
 
 import torch
diff --git a/distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py b/distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py
index 149de88..fb0f4bc 100644
--- a/distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py
+++ b/distributed_shampoo/preconditioner/tests/spectral_descent_preconditioner_list_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 from typing import Any
 
diff --git a/distributed_shampoo/shampoo_types.py b/distributed_shampoo/shampoo_types.py
index 98255fa..bd9cf7c 100644
--- a/distributed_shampoo/shampoo_types.py
+++ b/distributed_shampoo/shampoo_types.py
@@ -35,6 +35,7 @@
 # Keys for optimizer state (always checkpointed)
 FILTERED_GRAD = "filtered_grad"
 LR_SUM = "lr_sum"
+MOMENTUM_BUFFER = "momentum_buffer"
 STEP = "step"
 TRAIN_MODE = "train_mode"
 WEIGHT_BUFFER = "weight_buffer"
@@ -44,7 +45,6 @@
 BETAS = "betas"
 DISTRIBUTED_CONFIG = "distributed_config"
 EPSILON = "epsilon"
-EVAL_INTERP_COEFF = "eval_interp_coeff"
 GRAFTING_CONFIG = "grafting_config"
 ITERATE_AVERAGING_CONFIG = "iterate_averaging_config"
 LR = "lr"
@@ -54,14 +54,15 @@
 PRECONDITION_FREQUENCY = "precondition_frequency"
 PRECONDITIONER_CONFIG = "preconditioner_config"
 START_PRECONDITIONING_STEP = "start_preconditioning_step"
-TRAIN_INTERP_COEFF = "train_interp_coeff"
 USE_BIAS_CORRECTION = "use_bias_correction"
 USE_PIN_MEMORY = "use_pin_memory"
 WEIGHT_DECAY = "weight_decay"
 WEIGHT_DECAY_TYPE = "weight_decay_type"
 
 # Keys for lists of blocked states and metadata (never checkpointed)
+CUDA_STREAM = "cuda_stream"
 DISTRIBUTOR = "distributor"
+EVAL_INTERP_COEFF = "eval_interp_coeff"
 FILTERED_GRAD_LIST = "filtered_grad_list"
 GRAFTING_PRECONDITIONER_LIST = "grafting_preconditioner_list"
 LR_CPU_PINNED = "lr_cpu_pinned"
@@ -69,9 +70,12 @@
 MASKED_BLOCKED_GRADS = "masked_blocked_grads"
 MASKED_BLOCKED_PARAMS = "masked_blocked_params"
 MASKED_FILTERED_GRAD_LIST = "masked_filtered_grad_list"
+MASKED_MOMENTUM_BUFFER_LIST = "masked_momentum_buffer_list"
 MASKED_WEIGHT_BUFFER_LIST = "masked_weight_buffer_list"
+MOMENTUM_BUFFER_LIST = "momentum_buffer_list"
 PREVIOUS_GRAD_SELECTOR = "previous_grad_selector"
 SHAMPOO_PRECONDITIONER_LIST = "shampoo_preconditioner_list"
+TRAIN_INTERP_COEFF = "train_interp_coeff"
 WEIGHT_BUFFER_LIST = "weight_buffer_list"
 
 
@@ -154,28 +158,40 @@ class BaseShampooPreconditionerConfig(PreconditionerConfig):
 
     Attributes:
         amortized_computation_config (MatrixFunctionConfig): Configuration for the amortized computation, e.g., inverse-root computation or eigendecomposition.
-        num_tolerated_failed_amortized_computations (int): Number of failed amortized computations to tolerate before raising an error. (Default: 3)
-        factor_matrix_dtype (torch.dtype): Data type for factor matrix. (Default: torch.float32)
-        use_trace_scaling (bool): Flag for whether to normalize the factor matrix by its trace's sqrt before computing the inverse root.
-            Credit to https://arxiv.org/pdf/2506.03595. (Default: False)
         drop_weighting_factor_on_gsquare (bool): If True, drop the (1 - beta2) weighting factor when computing
             the updates of the preconditioners, i.e., V(t) = beta2 * V(t-1) + G^2 instead of
             V(t) = beta2 * V(t-1) + (1 - beta2) * G^2. This keeps _bias_correction2 at 1.0 (no beta2 bias correction). (Default: False)
+        factor_matrix_dtype (torch.dtype): Data type for factor matrix. (Default: torch.float32)
+        num_tolerated_failed_amortized_computations (int): Number of failed amortized computations to tolerate before raising an error. (Default: 3)
+        use_symmetric_packing (bool): If True, stores symmetric Kronecker factor matrices
+            (factor_matrices and inv_factor_matrices) in packed upper triangular format
+            (d*(d+1)/2 elements instead of d*d), saving ~50% memory. This is lossless since
+            these matrices are symmetric PSD. (Default: True)
+        use_trace_scaling (bool): Flag for whether to normalize the factor matrix by its trace raised to trace_scaling_exponent before computing the inverse root.
+            Credit to https://arxiv.org/pdf/2506.03595. (Default: False)
+        trace_scaling_exponent (float): Exponent applied to the trace when use_trace_scaling is True; the factor matrix is multiplied by trace ** (-trace_scaling_exponent).
+            Must be in (0.0, 1.0]. The default 0.5 reproduces 1 / sqrt(trace). (Default: 0.5)
 
     """
 
     # repr=False prevents __repr__() from accessing this field to avoid linter complaints
     amortized_computation_config: MatrixFunctionConfig = field(repr=False)
-    num_tolerated_failed_amortized_computations: int = 3
+    drop_weighting_factor_on_gsquare: bool = False
     factor_matrix_dtype: torch.dtype = torch.float32
+    num_tolerated_failed_amortized_computations: int = 3
+    use_symmetric_packing: bool = True
     use_trace_scaling: bool = False
-    drop_weighting_factor_on_gsquare: bool = False
+    trace_scaling_exponent: float = 0.5
 
     def __post_init__(self) -> None:
         if self.num_tolerated_failed_amortized_computations < 0:
             raise ValueError(
                 f"Invalid num_tolerated_failed_amortized_computations value: {self.num_tolerated_failed_amortized_computations}. Must be >= 0."
             )
+        if self.use_trace_scaling and not 0.0 < self.trace_scaling_exponent <= 1.0:
+            raise ValueError(
+                f"Invalid trace_scaling_exponent value: {self.trace_scaling_exponent}. Must be in (0.0, 1.0] when use_trace_scaling=True."
+            )
 
 
 @dataclass(init=False)
@@ -324,6 +340,7 @@ class RootInvShampooPreconditionerConfig(ClassicShampooPreconditionerConfig):
     def _get_default_amortized_computation_config() -> RootInvConfig:
         return DefaultEigenConfig
 
+    # pyrefly: ignore [bad-override-mutable-attribute]
     amortized_computation_config: RootInvConfig = field(
         default_factory=_get_default_amortized_computation_config
     )
@@ -390,6 +407,7 @@ class EigendecomposedShampooPreconditionerConfig(ClassicShampooPreconditionerCon
     def _get_default_amortized_computation_config() -> EigendecompositionConfig:
         return DefaultEigendecompositionConfig
 
+    # pyrefly: ignore [bad-override-mutable-attribute]
     amortized_computation_config: EigendecompositionConfig = field(
         default_factory=_get_default_amortized_computation_config
     )
@@ -458,6 +476,7 @@ class EigenvalueCorrectedShampooPreconditionerConfig(BaseShampooPreconditionerCo
     def _get_default_amortized_computation_config() -> EigendecompositionConfig:
         return DefaultEigendecompositionConfig
 
+    # pyrefly: ignore [bad-override-mutable-attribute]
     amortized_computation_config: EigendecompositionConfig = field(
         default_factory=_get_default_amortized_computation_config
     )
@@ -691,38 +710,11 @@ def _default_scale_fn(grad: Tensor) -> float:
 class IterateAveragingConfig(AbstractDataclass):
     """Configuration for primal or iterate averaging in Shampoo.
 
-    Iterate averaging methods like Generalized Primal Averaging (GPA) and Schedule-Free
-    provide an alternative to traditional momentum that can achieve equivalent or better
-    convergence properties.
-
-    Migration from Previous Momentum Parameters:
-        The previous `momentum`, `dampening`, and `use_nesterov` parameters have been replaced
-        by iterate averaging configs. Below are the equivalences:
-
-        **SGD Heavy-Ball Momentum (momentum=β, dampening=0, use_nesterov=False):**
-            Use GeneralizedPrimalAveragingConfig with:
-            - eval_interp_coeff = β (the momentum value)
-            - train_interp_coeff = 1.0
-            - Adjust learning rate: lr_new = lr_old / (1 - β)
-
-        **SGD Nesterov Momentum (momentum=β, dampening=0, use_nesterov=True):**
-            Use GeneralizedPrimalAveragingConfig with:
-            - eval_interp_coeff = β (the momentum value)
-            - train_interp_coeff = β (the momentum value)
-            - Adjust learning rate: lr_new = lr_old / (1 - β)
-
-        **Dampening (dampening ≠ 0):**
-            The previous dampening parameter does not have a direct equivalent in the
-            iterate averaging framework. If dampening was used, the behavior cannot be
-            exactly replicated. In practice, dampening was rarely used (default was 0.0),
-            and the primal averaging formulation provides better theoretical properties.
-
-        **LaProp:**
-            The previous momentum implementation (sometimes called LaProp in the codebase)
-            is mathematically equivalent to the heavy-ball/primal averaging formulation
-            when dampening=0. Use the heavy-ball configuration above.
-
-    Why Learning Rate Adjustment is Required:
+    Iterate averaging methods like Generalized Primal Averaging (GPA), Schedule-Free,
+    and ClassicMomentumConfig provide different approaches to incorporating momentum
+    into the optimizer.
+
+    Learning Rate Adjustment for Momentum -> GPA:
         In SGD heavy-ball momentum (dampening=0), the momentum buffer accumulates
         gradients and amplifies the effective step size by 1/(1-β) at steady state:
 
@@ -741,14 +733,24 @@ class IterateAveragingConfig(AbstractDataclass):
         step size is lr * (1-β) — a factor of 1/(1-β)^2 smaller than heavy-ball.
         To compensate, set lr_new = lr_old / (1-β).
 
-        This relationship is validated by `test_gpa_vs_sgd_momentum` in
+        ClassicMomentumConfig does NOT require this adjustment because it uses the
+        same momentum buffer accumulation as SGD.
+
+        These relationships are validated by `test_gpa_vs_sgd_momentum` and
+        `test_classic_momentum_vs_sgd_momentum` in
         `dev/gpu_tests/iterate_averaging_test.py`.
 
     Concrete Example:
-        An optimizer config with `lr=0.04, momentum=0.5` (heavy-ball) has an
-        effective step size of 0.04 / (1-0.5) = 0.08 due to momentum buffer
-        amplification. The equivalent GPA config is:
+        An optimizer config with `lr=0.04, momentum=0.5` (heavy-ball) can be implemented
+        two ways:
+
+        1. ClassicMomentumConfig (no lr change):
+            lr = 0.04
+            iterate_averaging_config = ClassicMomentumConfig(
+                momentum=0.5,
+            )
 
+        2. GPA (requires lr change):
             lr = 0.08  # = 0.04 / (1 - 0.5)
             iterate_averaging_config = GeneralizedPrimalAveragingConfig(
                 eval_interp_coeff=0.5,   # = momentum
@@ -758,6 +760,42 @@ class IterateAveragingConfig(AbstractDataclass):
     """
 
 
+@dataclass(kw_only=True)
+class ClassicMomentumConfig(IterateAveragingConfig):
+    """Configuration for classic SGD-style momentum in Shampoo.
+
+    This configures and enables the momentum/dampening/Nesterov behavior as traditional
+    SGD momentum applied to the preconditioned search direction, without requiring
+    train/eval mode switching (unlike GPA or Schedule-Free).
+
+    The momentum update applied to the search direction P is:
+        M <- momentum * M + (1 - dampening) * P
+        P <- (1 - dampening) * P + momentum * M     if use_nesterov
+        P <- M                                       otherwise
+
+    Then the parameter update is:
+        W <- W - lr * P
+
+    Attributes:
+        momentum (float): Momentum factor. Must be in (0.0, 1.0). (Default: 0.5)
+            Note: momentum=0.0 is a no-op and is rejected; pass
+            `iterate_averaging_config=None` instead to disable iterate averaging.
+        dampening (float): Dampening for momentum. Must be in [0.0, 1.0). (Default: 0.0)
+        use_nesterov (bool): Enables Nesterov momentum. (Default: False)
+
+    """
+
+    momentum: float = 0.5
+    dampening: float = 0.0
+    use_nesterov: bool = False
+
+    def __post_init__(self) -> None:
+        if not 0.0 < self.momentum < 1.0:
+            raise ValueError(f"Invalid {self.momentum=}. Must be within (0.0, 1.0).")
+        if not 0.0 <= self.dampening < 1.0:
+            raise ValueError(f"Invalid {self.dampening=}. Must be within [0.0, 1.0).")
+
+
 @dataclass(kw_only=True)
 class GeneralizedPrimalAveragingConfig(IterateAveragingConfig):
     """Configuration for generalized primal averaging in Shampoo.
@@ -1114,11 +1152,35 @@ class FullyShardDistributedConfig(DistributedConfig):
             (Default: 1)
         param_assignment_strategy (FSDPParamAssignmentStrategy): Strategy for assigning model parameters among the FSDP shards.
             (Default: FSDPParamAssignmentStrategy.DEFAULT)
+        num_sub_groups (int): Number of sub-groups to split a single param group into. Each sub-group gets its
+            own distributor and NCCL communicator, enabling multi-stream overlap of compute and communication.
+            Currently only effective with the *lossless* FSDP/HSDP distributors
+            (`_shampoo_fully_shard_lossless_distributor`,
+            `_shampoo_hybrid_shard_lossless_distributor`); other distributor
+            types ignore this field. (Default: 1)
     """
 
     param_assignment_strategy: FSDPParamAssignmentStrategy = (
         FSDPParamAssignmentStrategy.DEFAULT
     )
+    num_sub_groups: int = 1
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.num_sub_groups < 1:
+            raise ValueError(f"Invalid {self.num_sub_groups=}. Must be >= 1.")
+        # Sub-group splitting is only supported by the lossless distributors
+        # (selected when param_assignment_strategy is REPLICATE or ROUND_ROBIN).
+        if (
+            self.num_sub_groups > 1
+            and self.param_assignment_strategy == FSDPParamAssignmentStrategy.DEFAULT
+        ):
+            raise ValueError(
+                f"num_sub_groups > 1 requires param_assignment_strategy != DEFAULT "
+                f"(REPLICATE or ROUND_ROBIN); only the lossless FSDP/HSDP "
+                f"distributors support param-group splitting. Got "
+                f"{self.num_sub_groups=}, {self.param_assignment_strategy=}."
+            )
 
 
 @dataclass(kw_only=True)
@@ -1191,3 +1253,21 @@ class ShampooRuntimeConfig:
 
 
 DefaultShampooRuntimeConfig = ShampooRuntimeConfig()
+
+
+@dataclass(kw_only=True)
+class ConcurrencyConfig:
+    """
+    Controls concurrent execution of param-group steps.
+
+    Attributes:
+        enable_cuda_stream_for_param_groups (bool): Run each param group's GPU
+            work on its own CUDA stream so independent NCCL communicators can
+            overlap on-device. Requires all param groups to be on CUDA;
+            ignored on CPU. (Default: True)
+    """
+
+    enable_cuda_stream_for_param_groups: bool = True
+
+
+DefaultConcurrencyConfig = ConcurrencyConfig()
diff --git a/distributed_shampoo/tests/distributed_shampoo_test.py b/distributed_shampoo/tests/distributed_shampoo_test.py
index c2a95c7..c517749 100644
--- a/distributed_shampoo/tests/distributed_shampoo_test.py
+++ b/distributed_shampoo/tests/distributed_shampoo_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import abc
 import gc
 import logging
@@ -28,6 +26,7 @@
 )
 from distributed_shampoo.shampoo_types import (
     AdaGradPreconditionerConfig,
+    BaseShampooPreconditionerConfig,
     DefaultEigenvalueCorrectedShampooConfig,
     DefaultShampooConfig,
     DefaultSignDescentPreconditionerConfig,
@@ -49,6 +48,7 @@
     SpectralDescentPreconditionerConfig,
     WeightDecayType,
 )
+from distributed_shampoo.utils.shampoo_utils import pack_upper_triangular
 from torch import nn, Tensor
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
@@ -56,6 +56,18 @@
 )
 
 
+def _pack_if_enabled(
+    config: PreconditionerConfig, matrix: torch.Tensor
+) -> torch.Tensor:
+    """Pack matrix to upper triangular format if symmetric packing is enabled."""
+    return (
+        pack_upper_triangular(matrix)
+        if isinstance(config, BaseShampooPreconditionerConfig)
+        and config.use_symmetric_packing
+        else matrix
+    )
+
+
 @instantiate_parametrized_tests
 class DistributedShampooInitTest(unittest.TestCase):
     def setUp(self) -> None:
@@ -243,7 +255,7 @@ class NotSupportedDistributedConfig(DistributedConfig):
 
         self.assertRaisesRegex(
             NotImplementedError,
-            r"group\[DISTRIBUTED_CONFIG\]=.*\.NotSupportedDistributedConfig\(.*\) not supported!",
+            r"distributed_config=.*\.NotSupportedDistributedConfig\(.*\) not supported!",
             DistributedShampoo,
             params=self._model.parameters(),
             distributed_config=NotSupportedDistributedConfig(),
@@ -505,43 +517,19 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_0": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "inv_factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [1.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 1.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 1.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 1.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.eye(5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [1.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 1.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 1.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 1.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.eye(5)
                                 ),
                             },
                         },
@@ -567,43 +555,19 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_1": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "inv_factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [1.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 1.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 1.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 1.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.eye(5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [1.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 1.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 1.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 1.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 1.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.eye(5)
                                 ),
                             },
                         },
@@ -668,23 +632,11 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_0": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "factor_matrices_eigenvectors": {
@@ -734,23 +686,11 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_1": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "factor_matrices_eigenvectors": {
@@ -839,23 +779,11 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_0": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "factor_matrices_eigenvectors": {
@@ -910,23 +838,11 @@ def _ref_state_dict(self) -> dict[str, Any]:
                     "block_1": {
                         "shampoo": {
                             "factor_matrices": {
-                                0: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                0: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
-                                1: torch.tensor(
-                                    [
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        [0.0, 0.0, 0.0, 0.0, 0.0],
-                                    ]
+                                1: _pack_if_enabled(
+                                    self._preconditioner_config, torch.zeros(5, 5)
                                 ),
                             },
                             "factor_matrices_eigenvectors": {
@@ -1019,6 +935,54 @@ def _preconditioner_config(self) -> EigendecomposedKLShampooPreconditionerConfig
         return EigendecomposedKLShampooPreconditionerConfig()
 
 
+# ---- Unpacked (use_symmetric_packing=False) state dict test variants ----
+
+
+class ShampooStateDictUnpackedTest(ShampooStateDictTest):
+    @property
+    def _preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
+        return replace(DefaultShampooConfig, use_symmetric_packing=False)
+
+
+class EigendecomposedShampooStateDictUnpackedTest(EigendecomposedShampooStateDictTest):
+    @property
+    def _preconditioner_config(self) -> EigendecomposedShampooPreconditionerConfig:
+        return replace(
+            EigendecomposedShampooPreconditionerConfig(),
+            use_symmetric_packing=False,
+        )
+
+
+class EigenvalueCorrectedShampooStateDictUnpackedTest(
+    EigenvalueCorrectedShampooStateDictTest,
+):
+    @property
+    def _preconditioner_config(self) -> EigenvalueCorrectedShampooPreconditionerConfig:
+        return replace(
+            DefaultEigenvalueCorrectedShampooConfig,
+            use_symmetric_packing=False,
+        )
+
+
+class RootInvKLShampooStateDictUnpackedTest(RootInvKLShampooStateDictTest):
+    @property
+    def _preconditioner_config(self) -> RootInvKLShampooPreconditionerConfig:
+        return replace(
+            RootInvKLShampooPreconditionerConfig(), use_symmetric_packing=False
+        )
+
+
+class EigendecomposedKLShampooStateDictUnpackedTest(
+    EigendecomposedKLShampooStateDictTest,
+):
+    @property
+    def _preconditioner_config(self) -> EigendecomposedKLShampooPreconditionerConfig:
+        return replace(
+            EigendecomposedKLShampooPreconditionerConfig(),
+            use_symmetric_packing=False,
+        )
+
+
 class SignDescentStateDictTest(AbstractTest.NoPreconditionerStateDictTestBase):
     @property
     def _preconditioner_config(self) -> SignDescentPreconditionerConfig:
@@ -1148,43 +1112,19 @@ def _ref_state_dict(self) -> dict[str, Any]:
                         "block_0": {
                             "shampoo": {
                                 "factor_matrices": {
-                                    0: torch.tensor(
-                                        [
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        ]
+                                    0: _pack_if_enabled(
+                                        self._preconditioner_config, torch.zeros(5, 5)
                                     ),
-                                    1: torch.tensor(
-                                        [
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        ]
+                                    1: _pack_if_enabled(
+                                        self._preconditioner_config, torch.zeros(5, 5)
                                     ),
                                 },
                                 "inv_factor_matrices": {
-                                    0: torch.tensor(
-                                        [
-                                            [1.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 1.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 1.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 1.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 1.0],
-                                        ]
+                                    0: _pack_if_enabled(
+                                        self._preconditioner_config, torch.eye(5)
                                     ),
-                                    1: torch.tensor(
-                                        [
-                                            [1.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 1.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 1.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 1.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 1.0],
-                                        ]
+                                    1: _pack_if_enabled(
+                                        self._preconditioner_config, torch.eye(5)
                                     ),
                                 },
                             },
@@ -1219,43 +1159,19 @@ def _ref_state_dict(self) -> dict[str, Any]:
                         "block_1": {
                             "shampoo": {
                                 "factor_matrices": {
-                                    0: torch.tensor(
-                                        [
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        ]
+                                    0: _pack_if_enabled(
+                                        self._preconditioner_config, torch.zeros(5, 5)
                                     ),
-                                    1: torch.tensor(
-                                        [
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 0.0],
-                                        ]
+                                    1: _pack_if_enabled(
+                                        self._preconditioner_config, torch.zeros(5, 5)
                                     ),
                                 },
                                 "inv_factor_matrices": {
-                                    0: torch.tensor(
-                                        [
-                                            [1.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 1.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 1.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 1.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 1.0],
-                                        ]
+                                    0: _pack_if_enabled(
+                                        self._preconditioner_config, torch.eye(5)
                                     ),
-                                    1: torch.tensor(
-                                        [
-                                            [1.0, 0.0, 0.0, 0.0, 0.0],
-                                            [0.0, 1.0, 0.0, 0.0, 0.0],
-                                            [0.0, 0.0, 1.0, 0.0, 0.0],
-                                            [0.0, 0.0, 0.0, 1.0, 0.0],
-                                            [0.0, 0.0, 0.0, 0.0, 1.0],
-                                        ]
+                                    1: _pack_if_enabled(
+                                        self._preconditioner_config, torch.eye(5)
                                     ),
                                 },
                             },
@@ -1390,6 +1306,18 @@ def _iterate_averaging_config(self) -> ScheduleFreeConfig:
         )
 
 
+class GPAShampooStateDictUnpackedTest(GPAShampooStateDictTest):
+    @property
+    def _preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
+        return replace(DefaultShampooConfig, use_symmetric_packing=False)
+
+
+class ScheduleFreeShampooStateDictUnpackedTest(ScheduleFreeShampooStateDictTest):
+    @property
+    def _preconditioner_config(self) -> RootInvShampooPreconditionerConfig:
+        return replace(DefaultShampooConfig, use_symmetric_packing=False)
+
+
 @instantiate_parametrized_tests
 class DistributedShampooTrainEvalModeTest(unittest.TestCase):
     """Test train/eval mode switching with iterate averaging configurations."""
diff --git a/distributed_shampoo/tests/shampoo_test_utils.py b/distributed_shampoo/tests/shampoo_test_utils.py
index 5572ab3..7a1c264 100644
--- a/distributed_shampoo/tests/shampoo_test_utils.py
+++ b/distributed_shampoo/tests/shampoo_test_utils.py
@@ -16,9 +16,37 @@
 from torch import nn
 from torch.distributed.fsdp import FSDPModule, fully_shard, FullyShardedDataParallel
 from torch.distributed.tensor import DTensor
+from torch.nn.parallel import DistributedDataParallel
 from torch.optim.optimizer import ParamsT
 
 
+def generate_global_train_data(
+    num_steps: int,
+    world_size: int,
+    data_shape: tuple[int, ...],
+    device: torch.device | None = None,
+    seed: int = 42,
+) -> torch.Tensor:
+    """Generate global training data with shape (num_steps, world_size, *data_shape).
+
+    Args:
+        num_steps (int): Number of training steps.
+        world_size (int): Number of data-parallel workers.
+        data_shape (tuple[int, ...]): Shape of each data sample.
+        device (torch.device | None): Device to place data on. (Default: None)
+        seed (int): Random seed for reproducibility. (Default: 42)
+
+    Returns:
+        torch.Tensor: Global training data of shape (num_steps, world_size, *data_shape).
+    """
+    torch.manual_seed(seed)
+    data = torch.randn(num_steps, world_size, *data_shape, dtype=torch.float)
+    data /= torch.linalg.norm(data, dim=-1, keepdim=True)
+    if device is not None:
+        data = data.to(device=device)
+    return data
+
+
 class _ModelWithScalarAndLinearAndDeadLayers(nn.Module):
     def __init__(
         self,
@@ -100,7 +128,6 @@ def construct_training_problem(
         fill (float | tuple[float, ...]): The value(s) to fill the model parameters. If a tuple, each element should correspond to one layer. (Default: 0.0)
         post_model_decoration (Callable[[nn.Module], nn.Module | FSDPModule]): A function to apply additional modifications to the model, useful for FullyShardedDataParallel and FSDPModule. (Default: identity function)
 
-
     Returns:
         model (nn.Module | FSDPModule): The model as specified from the input arguments.
         loss (nn.Module): The loss function (currently always set to MSE).
@@ -136,7 +163,8 @@ def construct_training_problem(
 
     loss = nn.MSELoss()
 
-    target = torch.tensor([[0.0] * model_linear_layers_dims[-1]]).to(device=device)
+    # Citrine C3: create tensor directly on device
+    target = torch.tensor([[0.0] * model_linear_layers_dims[-1]], device=device)
 
     return post_model_decoration(model), loss, data, target
 
@@ -148,6 +176,7 @@ def train_model(
         [], tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]
     ],
     num_steps: int = 5,
+    train_data: torch.Tensor | None = None,
 ) -> tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor, torch.optim.Optimizer]: ...
 
 
@@ -158,6 +187,7 @@ def train_model(
         [], tuple[FSDPModule, nn.Module, torch.Tensor, torch.Tensor]
     ],
     num_steps: int = 5,
+    train_data: torch.Tensor | None = None,
 ) -> tuple[
     FSDPModule, nn.Module, torch.Tensor, torch.Tensor, torch.optim.Optimizer
 ]: ...
@@ -169,6 +199,7 @@ def train_model(
         [], tuple[nn.Module | FSDPModule, nn.Module, torch.Tensor, torch.Tensor]
     ],
     num_steps: int = 5,
+    train_data: torch.Tensor | None = None,
 ) -> tuple[
     nn.Module | FSDPModule, nn.Module, torch.Tensor, torch.Tensor, torch.optim.Optimizer
 ]:
@@ -183,6 +214,9 @@ def train_model(
         optim_factory (Callable[[ParamsT], torch.optim.Optimizer]): A factory function that returns an optimizer instance.
         model_factory (Callable[[], tuple[nn.Module | FSDPModule, nn.Module, torch.Tensor, torch.Tensor]]): A factory function that returns a tuple containing the model, loss function, validation data, and target data.
         num_steps (int): The number of training steps to perform. (Default: 5)
+        train_data (torch.Tensor | None): Pre-generated training data. If provided, num_steps is derived from train_data.shape[0].
+            Shape can be (num_steps, *data_shape) for single-sample-per-step or (num_steps, world_size, *data_shape) for
+            full-batch training. If None, generates data internally with seed 42. (Default: None)
 
     Returns:
         model (nn.Module | FSDPModule): The trained model.
@@ -196,17 +230,18 @@ def train_model(
     assert isinstance(model, nn.Module)
     optimizer = optim_factory(model.parameters())
 
-    # Pregenerate num_steps of tensor and put into a PyTorch dataset with seed
-    seed_value = 42
-    torch.manual_seed(seed_value)
-
-    train_data = torch.randn(num_steps, *validation_data.shape, dtype=torch.float)
-    train_data /= torch.linalg.norm(train_data, dim=1, keepdim=True)
+    if train_data is None:
+        # Generate training data internally (backward compat for single-device tests).
+        seed_value = 42
+        torch.manual_seed(seed_value)
+        train_data = torch.randn(num_steps, *validation_data.shape, dtype=torch.float)
+        train_data /= torch.linalg.norm(train_data, dim=-1, keepdim=True)
+
+    # Move data onto the model's device. Pre-generated global data (e.g. the CPU-vs-device
+    # comparisons) is created once on a single device for value-determinism across devices,
+    # so the control model on CPU must relocate its data here to avoid a device mismatch.
+    # pyrefly: ignore [missing-attribute]
     train_data = train_data.to(device=validation_data.device)
-    train_dataset = torch.utils.data.TensorDataset(train_data)
-    train_loader = torch.utils.data.DataLoader(
-        train_dataset, batch_size=1, shuffle=False
-    )
 
     # Set optimizer to train mode if it supports train/eval modes (e.g., GPA-AdamW).
     train_mode = getattr(optimizer, "train", None)
@@ -214,10 +249,26 @@ def train_model(
         train_mode()
 
     # Train for the specified number of steps
-    for _, (step_data,) in enumerate(train_loader):
+    for step in range(train_data.shape[0]):
         optimizer.zero_grad()
-        objective = loss(model(step_data), target)
-        objective.backward()
+        step_data = train_data[step]
+
+        if step_data.ndim > validation_data.ndim:
+            # Full-batch data with an extra leading batch dimension
+            # (e.g., shape (world_size, *data_shape) vs validation_data shape (*data_shape)).
+            # Process each sample individually and average gradients to match
+            # DDP's per-sample computation + all-reduce averaging exactly.
+            batch_size = step_data.shape[0]
+            for i in range(batch_size):
+                objective = loss(model(step_data[i]), target)
+                objective.backward()
+            for param in model.parameters():
+                if param.grad is not None:
+                    param.grad /= batch_size
+        else:
+            objective = loss(model(step_data), target)
+            objective.backward()
+
         optimizer.step()
 
     return model, loss, validation_data.unsqueeze(0), target, optimizer
@@ -236,6 +287,8 @@ def compare_two_optimizers_models_devices_on_weight_and_loss(
     total_steps: int = 5,
     rtol: float | None = None,
     atol: float | None = None,
+    control_train_data: torch.Tensor | None = None,
+    experimental_train_data: torch.Tensor | None = None,
 ) -> None:
     """
     Compare the performance of two optimizers on models across different devices by evaluating their weights and loss.
@@ -249,6 +302,8 @@ def compare_two_optimizers_models_devices_on_weight_and_loss(
         total_steps (int): Number of training steps. (Default: 5)
         rtol (float | None): Relative tolerance for comparing weights and losses. (Default: None)
         atol (float | None): Absolute tolerance for comparing weights and losses. (Default: None)
+        control_train_data (torch.Tensor | None): Pre-generated training data for the control optimizer. (Default: None)
+        experimental_train_data (torch.Tensor | None): Pre-generated training data for the experimental optimizer. (Default: None)
 
     Returns:
         None
@@ -259,6 +314,7 @@ def generated_comparable_trained_weights_and_final_loss(
         model_factory: Callable[
             [], tuple[nn.Module | FSDPModule, nn.Module, torch.Tensor, torch.Tensor]
         ],
+        train_data: torch.Tensor | None = None,
     ) -> tuple[list[torch.Tensor], torch.Tensor]:
         """
         Trains a model and returns the trained weights and final loss.
@@ -266,6 +322,7 @@ def generated_comparable_trained_weights_and_final_loss(
         Args:
             optim_factory (Callable[[ParamsT], torch.optim.Optimizer]): A factory function to create an optimizer.
             model_factory (Callable[[], tuple[nn.Module | FSDPModule, nn.Module, torch.Tensor, torch.Tensor]]): A factory function to create a model, loss, data, and target.
+            train_data (torch.Tensor | None): Pre-generated training data. (Default: None)
 
         Returns:
             trained_weights (list[torch.Tensor]): A list of trained model parameters.
@@ -274,27 +331,35 @@ def generated_comparable_trained_weights_and_final_loss(
         # Using partial here to prevent Pyre complaints on incompatible parameter type.
         model, loss, validation_data, target, _ = partial(
             train_model, model_factory=model_factory
-        )(optim_factory=optim_factory, num_steps=total_steps)
+        )(optim_factory=optim_factory, num_steps=total_steps, train_data=train_data)
 
         # We only care about model_linear_layers_dim params, not model_dead_layer params.
         assert isinstance(model, nn.Module)
-        linear_layers = model.get_submodule("linear_layers")
         match model:
+            case DistributedDataParallel():
+                linear_layers = model.module.get_submodule("linear_layers")
+                trained_weights = [
+                    param.view(-1).detach().cpu()
+                    for param in linear_layers.parameters()
+                ]
             case FSDPModule():
                 # When FullyShard or Hybrid Shard is used, model parameters are DTensors. We obtain the full value of
                 # parameters from DTensors.
+                linear_layers = model.get_submodule("linear_layers")
                 trained_weights = []
                 for param in linear_layers.parameters():
                     # Need this assertion to pass the type-checking test.
                     assert isinstance(param, DTensor)
                     trained_weights.append(param.full_tensor().view(-1).detach().cpu())
             case FullyShardedDataParallel():
+                linear_layers = model.get_submodule("linear_layers")
                 with FullyShardedDataParallel.summon_full_params(model):
                     trained_weights = [
                         param.view(-1).detach().cpu()
                         for param in linear_layers.parameters()
                     ]
             case _:
+                linear_layers = model.get_submodule("linear_layers")
                 trained_weights = [
                     param.view(-1).detach().cpu()
                     for param in linear_layers.parameters()
@@ -305,11 +370,13 @@ def generated_comparable_trained_weights_and_final_loss(
     control_params, control_loss = generated_comparable_trained_weights_and_final_loss(
         optim_factory=control_optim_factory,
         model_factory=control_model_factory,
+        train_data=control_train_data,
     )
     experimental_params, experimental_loss = (
         generated_comparable_trained_weights_and_final_loss(
             optim_factory=experimental_optim_factory,
             model_factory=experimental_model_factory,
+            train_data=experimental_train_data,
         )
     )
     torch.testing.assert_close(
@@ -340,6 +407,9 @@ def compare_two_optimizers_devices_on_weight_and_loss(
     total_steps: int = 5,
     rtol: float | None = None,
     atol: float | None = None,
+    rank: int = 0,
+    world_size: int = 1,
+    experimental_post_model_decoration: Callable[[nn.Module], nn.Module] = lambda x: x,
 ) -> None:
     """
     Compare the performance of two optimizers on a simple neural network across different devices.
@@ -356,6 +426,11 @@ def compare_two_optimizers_devices_on_weight_and_loss(
         total_steps (int): The number of training steps. (Default: 5)
         rtol (float | None): The relative tolerance for comparing weights and losses. (Default: None)
         atol (float | None): The absolute tolerance for comparing weights and losses. (Default: None)
+        rank (int): The rank of the current process. (Default: 0)
+        world_size (int): The number of data-parallel workers. When world_size > 1, the control optimizer receives
+            the full global batch while the experimental optimizer receives its rank's shard. (Default: 1)
+        experimental_post_model_decoration (Callable[[nn.Module], nn.Module]): A function to apply additional modifications
+            to the experimental model (e.g., DDP wrapping). (Default: identity function)
 
     Returns:
         None
@@ -368,6 +443,27 @@ def compare_two_optimizers_devices_on_weight_and_loss(
         enable_learnable_scalar=enable_learnable_scalar,
         fill=fill,
     )
+
+    # Generate global training data for proper data-parallel sharding.
+    input_dim = model_linear_layers_dims[0]
+    # Use the experimental device for data generation (both devices see the same data).
+    data_device = experimental_device or control_device
+    global_train_data = generate_global_train_data(
+        num_steps=total_steps,
+        world_size=world_size,
+        data_shape=(input_dim,),
+        device=data_device,
+    )
+
+    if world_size == 1:
+        # Single device: both sides get the same (total_steps, input_dim) data.
+        control_train_data = global_train_data.squeeze(1)
+        experimental_train_data = global_train_data.squeeze(1)
+    else:
+        # Data parallel: control gets full batch, experimental gets its shard.
+        control_train_data = global_train_data
+        experimental_train_data = global_train_data[:, rank]
+
     compare_two_optimizers_models_devices_on_weight_and_loss(
         control_optim_factory=control_optim_factory,
         control_model_factory=partial(
@@ -375,12 +471,16 @@ def compare_two_optimizers_devices_on_weight_and_loss(
         ),
         experimental_optim_factory=experimental_optim_factory,
         experimental_model_factory=partial(
-            device_aware_training_problem_factory, device=experimental_device
+            device_aware_training_problem_factory,
+            device=experimental_device,
+            post_model_decoration=experimental_post_model_decoration,
         ),
         control_and_experimental_same_device=control_device == experimental_device,
         total_steps=total_steps,
         rtol=rtol,
         atol=atol,
+        control_train_data=control_train_data,
+        experimental_train_data=experimental_train_data,
     )
 
 
@@ -395,6 +495,9 @@ def compare_two_optimizers_on_weight_and_loss(
     total_steps: int = 5,
     rtol: float | None = None,
     atol: float | None = None,
+    rank: int = 0,
+    world_size: int = 1,
+    experimental_post_model_decoration: Callable[[nn.Module], nn.Module] = lambda x: x,
 ) -> None:
     """
     Compare the performance of two optimizers on a simple neural network using the same device.
@@ -410,6 +513,10 @@ def compare_two_optimizers_on_weight_and_loss(
         total_steps (int): The number of training steps. (Default: 5)
         rtol (float | None): The relative tolerance for comparing weights and losses. (Default: None)
         atol (float | None): The absolute tolerance for comparing weights and losses. (Default: None)
+        rank (int): The rank of the current process. (Default: 0)
+        world_size (int): The number of data-parallel workers. (Default: 1)
+        experimental_post_model_decoration (Callable[[nn.Module], nn.Module]): A function to apply additional modifications
+            to the experimental model (e.g., DDP wrapping). (Default: identity function)
 
     Returns:
         None
@@ -426,6 +533,9 @@ def compare_two_optimizers_on_weight_and_loss(
         total_steps=total_steps,
         rtol=rtol,
         atol=atol,
+        rank=rank,
+        world_size=world_size,
+        experimental_post_model_decoration=experimental_post_model_decoration,
     )
 
 
@@ -439,6 +549,8 @@ def compare_optimizer_on_cpu_and_device(
     total_steps: int = 5,
     rtol: float | None = None,
     atol: float | None = None,
+    rank: int = 0,
+    world_size: int = 1,
 ) -> None:
     """
     Compare the performance of the same optimizer on a simple neural network across CPU and another device.
@@ -453,6 +565,8 @@ def compare_optimizer_on_cpu_and_device(
         total_steps (int): The number of training steps. (Default: 5)
         rtol (float | None): The relative tolerance for comparing weights and losses. (Default: None)
         atol (float | None): The absolute tolerance for comparing weights and losses. (Default: None)
+        rank (int): The rank of the current process. (Default: 0)
+        world_size (int): The number of data-parallel workers. (Default: 1)
 
     Returns:
         None
@@ -469,4 +583,6 @@ def compare_optimizer_on_cpu_and_device(
         total_steps=total_steps,
         rtol=rtol,
         atol=atol,
+        rank=rank,
+        world_size=world_size,
     )
diff --git a/distributed_shampoo/tests/shampoo_types_test.py b/distributed_shampoo/tests/shampoo_types_test.py
index 9fd368a..c255e87 100644
--- a/distributed_shampoo/tests/shampoo_types_test.py
+++ b/distributed_shampoo/tests/shampoo_types_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 from typing import Any
@@ -22,16 +20,17 @@
 from distributed_shampoo.shampoo_types import (
     AdaGradPreconditionerConfig,
     BaseShampooPreconditionerConfig,
+    ClassicMomentumConfig,
     ClassicShampooPreconditionerConfig,
     DistributedConfig,
     EigenvalueCorrectedShampooPreconditionerConfig,
     FSDPDistributedConfig,
+    FullyShardDistributedConfig,
     GeneralizedPrimalAveragingConfig,
     HSDPDistributedConfig,
     HybridShardDistributedConfig,
     IterateAveragingConfig,
     RMSpropPreconditionerConfig,
-    ScheduleFreeConfig,
     SignDescentPreconditionerConfig,
 )
 from distributed_shampoo.utils.commons import get_all_non_abstract_subclasses
@@ -299,9 +298,12 @@ def l1_norm_scale_fn(grad: torch.Tensor) -> float:
 
 @instantiate_parametrized_tests
 class IterateAveragingConfigSubclassesTest(unittest.TestCase):
-    subclasses_types: list[type[IterateAveragingConfig]] = list(
-        get_all_non_abstract_subclasses(IterateAveragingConfig)  # type: ignore[type-abstract]
-    )
+    # Only test subclasses that have a train_interp_coeff field.
+    subclasses_types: list[type[IterateAveragingConfig]] = [
+        cls  # type: ignore[type-abstract]
+        for cls in get_all_non_abstract_subclasses(IterateAveragingConfig)
+        if hasattr(cls, "train_interp_coeff")
+    ]
 
     @parametrize("train_interp_coeff", (-0.1, 0.0, 1.5))
     @parametrize("cls", subclasses_types)
@@ -331,23 +333,28 @@ def test_illegal_eval_interp_coeff(self, eval_interp_coeff: float) -> None:
             eval_interp_coeff=eval_interp_coeff,
         )
 
-    def test_valid_config(self) -> None:
-        # Test that valid configurations do not raise exceptions.
-        # Boundary values for eval_interp_coeff: 0.0 (inclusive) and just below 1.0.
-        # Boundary values for train_interp_coeff: just above 0.0 (exclusive) and 1.0 (inclusive).
-        GeneralizedPrimalAveragingConfig(eval_interp_coeff=0.0, train_interp_coeff=0.01)
-        GeneralizedPrimalAveragingConfig(eval_interp_coeff=0.5, train_interp_coeff=0.5)
-        GeneralizedPrimalAveragingConfig(eval_interp_coeff=0.99, train_interp_coeff=1.0)
-
 
 @instantiate_parametrized_tests
-class ScheduleFreeConfigTest(unittest.TestCase):
-    def test_valid_config(self) -> None:
-        # Test that valid configurations do not raise exceptions.
-        # Boundary values: just above 0.0 (exclusive) and 1.0 (inclusive).
-        ScheduleFreeConfig(train_interp_coeff=0.01)
-        ScheduleFreeConfig(train_interp_coeff=0.5)
-        ScheduleFreeConfig(train_interp_coeff=1.0)
+class ClassicMomentumConfigTest(unittest.TestCase):
+    @parametrize("momentum", (-0.1, 0.0, 1.0, 1.5))
+    def test_illegal_momentum(self, momentum: float) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(f"Invalid self.momentum={momentum}. Must be within (0.0, 1.0)."),
+            ClassicMomentumConfig,
+            momentum=momentum,
+        )
+
+    @parametrize("dampening", (-0.1, 1.0, 1.5))
+    def test_illegal_dampening(self, dampening: float) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(
+                f"Invalid self.dampening={dampening}. Must be within [0.0, 1.0)."
+            ),
+            ClassicMomentumConfig,
+            dampening=dampening,
+        )
 
 
 @instantiate_parametrized_tests
@@ -389,3 +396,28 @@ def test_illegal_target_parameter_dimensionality(
             cls,
             **kwargs,
         )
+
+
+@instantiate_parametrized_tests
+class FullyShardDistributedConfigTest(unittest.TestCase):
+    @parametrize("num_sub_groups", (0, -1))
+    def test_illegal_num_sub_groups(self, num_sub_groups: int) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(f"Invalid self.num_sub_groups={num_sub_groups}. Must be >= 1."),
+            FullyShardDistributedConfig,
+            num_sub_groups=num_sub_groups,
+        )
+
+
+@instantiate_parametrized_tests
+class HybridShardDistributedConfigTest(unittest.TestCase):
+    @parametrize("num_sub_groups", (0, -1))
+    def test_illegal_num_sub_groups(self, num_sub_groups: int) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(f"Invalid self.num_sub_groups={num_sub_groups}. Must be >= 1."),
+            HybridShardDistributedConfig,
+            device_mesh=MagicMock(),
+            num_sub_groups=num_sub_groups,
+        )
diff --git a/distributed_shampoo/utils/__init__.py b/distributed_shampoo/utils/__init__.py
index e69de29..9071dd2 100644
--- a/distributed_shampoo/utils/__init__.py
+++ b/distributed_shampoo/utils/__init__.py
@@ -0,0 +1,8 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
diff --git a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
index 9884a2a..994eb55 100644
--- a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
+++ b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
@@ -7,15 +7,17 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 
 import numpy as np
 import torch
-from distributed_shampoo.utils.shampoo_utils import (
+from distributed_shampoo.utils.shampoo_fully_shard_utils import (
+    _compute_chunk_sizes,
+    _compute_param_chunk_sizes,
+    GatherGradientsContext,
     prepare_update_param_buffers,
     redistribute_and_update_params,
+    RedistributeParamsContext,
 )
 from torch import distributed as dist
 from torch.distributed.device_mesh import init_device_mesh
@@ -84,3 +86,754 @@ def test_redistribute_and_update_params(self, num_params: int) -> None:
             np.testing.assert_allclose(
                 param.to_local().cpu().numpy(), i % self.world_size
             )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+@instantiate_parametrized_tests
+class RedistributeParamsContextTest(DTensorTestBase):
+    """Tests for RedistributeParamsContext class.
+
+    Validates initialization, metadata precomputation, and parameter redistribution
+    using a single coalesced all_to_all collective.
+    """
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_init_empty_params_raises_value_error(self) -> None:
+        """Verify that empty params raises ValueError."""
+        dist_group = dist.distributed_c10d._get_default_group()
+        with self.assertRaises(ValueError):
+            RedistributeParamsContext(
+                params=(),
+                assigned_params_mask=(),
+                dist_group=dist_group,
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_init_mismatched_dtypes_raises_not_implemented(self) -> None:
+        """Verify that params with different dtypes raises NotImplementedError."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        param_f32 = distribute_tensor(
+            torch.zeros(4, 4, device="cuda", dtype=torch.float32),
+            device_mesh,
+            [Shard(0)],
+        )
+        param_f64 = distribute_tensor(
+            torch.zeros(4, 4, device="cuda", dtype=torch.float64),
+            device_mesh,
+            [Shard(0)],
+        )
+        dist_group = dist.distributed_c10d._get_default_group()
+        with self.assertRaises(NotImplementedError):
+            RedistributeParamsContext(
+                params=(param_f32, param_f64),
+                assigned_params_mask=(True, False),
+                dist_group=dist_group,
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_precomputed_metadata_chunk_sizes(self) -> None:
+        """Verify that chunk sizes are correctly precomputed for different param shapes."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        # Create params with dim-0 size of 8 (evenly divisible by 4 ranks)
+        param1 = distribute_tensor(
+            torch.zeros(8, 3, device="cuda"), device_mesh, [Shard(0)]
+        )
+        # Create params with dim-0 size of 5 (not evenly divisible by 4 ranks)
+        param2 = distribute_tensor(
+            torch.zeros(5, 2, device="cuda"), device_mesh, [Shard(0)]
+        )
+        dist_group = dist.distributed_c10d._get_default_group()
+        ctx = RedistributeParamsContext(
+            params=(param1, param2),
+            assigned_params_mask=(True, False),
+            dist_group=dist_group,
+        )
+
+        # param1: shape (8, 3), dim-0 chunked into 4 → [2, 2, 2, 2] * 3 = [6, 6, 6, 6]
+        self.assertEqual(ctx._param_chunk_sizes[0], [6, 6, 6, 6])
+        # param2: shape (5, 2), dim-0 chunked into 4 → [2, 2, 1, 0] * 2 = [4, 4, 2, 0]
+        # (torch.chunk(5, 4) → ceil(5/4)=2, chunks of [2,2,1])
+        # Actually: ceil(5/4) = 2, so chunks = [2, 2, 1, 0] → *2 = [4, 4, 2, 0]
+        expected_dim0_chunks = _compute_chunk_sizes(5, 4)
+        expected_chunk_sizes_param2 = [s * 2 for s in expected_dim0_chunks]
+        self.assertEqual(ctx._param_chunk_sizes[1], expected_chunk_sizes_param2)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_precomputed_metadata_send_recv_sizes(self) -> None:
+        """Verify send and recv sizes are correctly computed."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        shapes = generate_param_shapes(4)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+        # Round-robin: rank 0 owns param 0, rank 1 owns param 1, etc.
+        assigned_mask = tuple(i % 4 == rank for i in range(4))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        # Verify send_sizes sum equals total_send_size
+        self.assertEqual(sum(ctx._send_sizes), ctx._total_send_size)
+        # Verify recv_sizes sum equals total_recv_size
+        self.assertEqual(sum(ctx._recv_sizes), ctx._total_recv_size)
+        # Verify recv buffer is pre-allocated with correct size
+        self.assertEqual(ctx._recv_buffer.numel(), ctx._total_recv_size)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_precomputed_metadata_local_param_indices(self) -> None:
+        """Verify local param indices are correctly identified from assigned_params_mask."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        shapes = generate_param_shapes(7)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+        # Round-robin assignment
+        assigned_mask = tuple(i % 4 == rank for i in range(7))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        expected_indices = [i for i in range(7) if i % 4 == rank]
+        self.assertEqual(ctx._local_param_indices, expected_indices)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize("num_params", (1, 4, 7))
+    def test_redistribute_params_context_correctness(self, num_params: int) -> None:
+        """Verify RedistributeParamsContext produces correct results.
+
+        Each rank fills its assigned params with its rank value, then redistributes.
+        After redistribution, each param's local shard should contain the value of the
+        rank that owns that param (param i is owned by rank i % world_size).
+        """
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        shapes = generate_param_shapes(num_params)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+
+        # Round-robin assignment
+        assigned_mask = tuple(i % self.world_size == rank for i in range(num_params))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        # Fill locally assigned parameters with rank value
+        local_full_params = [
+            torch.zeros(s, device="cuda").fill_(rank)
+            for i, s in enumerate(shapes)
+            if i % self.world_size == rank
+        ]
+
+        ctx.redistribute_and_update_params(local_full_params)
+
+        # Verify each param's local shard has the correct value
+        for i, param in enumerate(dtensor_params):
+            owning_rank = i % self.world_size
+            local_data = param.to_local().cpu().numpy()
+            np.testing.assert_allclose(
+                local_data,
+                owning_rank,
+                err_msg=f"Param {i} (owned by rank {owning_rank}) has incorrect values "
+                f"on rank {rank}",
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_redistribute_params_context_wrong_num_local_params_asserts(
+        self,
+    ) -> None:
+        """Verify assertion when local_full_params count doesn't match."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        shapes = generate_param_shapes(4)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+        assigned_mask = tuple(i % 4 == rank for i in range(4))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        # Provide wrong number of local params
+        wrong_local_params = [
+            torch.zeros(1, device="cuda"),
+            torch.zeros(1, device="cuda"),
+        ]
+        with self.assertRaises(AssertionError):
+            ctx.redistribute_and_update_params(wrong_local_params)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_redistribute_params_context_matches_legacy(self) -> None:
+        """Verify RedistributeParamsContext produces identical results to the legacy
+        redistribute_and_update_params function."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        num_params = 7
+        shapes = generate_param_shapes(num_params)
+
+        # Create two sets of identical DTensor params
+        params_legacy = [torch.zeros(s, device="cuda") for s in shapes]
+        params_new = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params_legacy = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params_legacy
+        )
+        dtensor_params_new = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params_new
+        )
+
+        dist_group = dist.distributed_c10d._get_default_group()
+        assigned_mask = tuple(i % self.world_size == rank for i in range(num_params))
+
+        # Create the same local_full_params for both paths
+        # Use distinct values per param to verify correctness
+        local_full_params_legacy = [
+            torch.zeros(s, device="cuda").fill_(float(rank * 100 + i))
+            for i, s in enumerate(shapes)
+            if i % self.world_size == rank
+        ]
+        local_full_params_new = [
+            torch.zeros(s, device="cuda").fill_(float(rank * 100 + i))
+            for i, s in enumerate(shapes)
+            if i % self.world_size == rank
+        ]
+
+        # Run legacy path
+        update_buffers = prepare_update_param_buffers(
+            dtensor_params_legacy, self.world_size
+        )
+        redistribute_and_update_params(
+            dtensor_params_legacy,
+            local_full_params_legacy,
+            update_buffers,
+            dist_group,
+        )
+
+        # Run new path
+        ctx = RedistributeParamsContext(
+            params=dtensor_params_new,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+        ctx.redistribute_and_update_params(local_full_params_new)
+
+        # Compare results
+        for i in range(num_params):
+            legacy_vals = dtensor_params_legacy[i].to_local().cpu().numpy()
+            new_vals = dtensor_params_new[i].to_local().cpu().numpy()
+            np.testing.assert_allclose(
+                new_vals,
+                legacy_vals,
+                err_msg=f"Param {i} mismatch between legacy and new redistribution on rank {rank}",
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_recv_buffer_reuse(self) -> None:
+        """Verify that the recv buffer is reused across multiple redistribute calls."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        shapes = generate_param_shapes(4)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+        assigned_mask = tuple(i % 4 == rank for i in range(4))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        # Capture the recv buffer's data_ptr
+        recv_buffer_ptr = ctx._recv_buffer.data_ptr()
+
+        # Run redistribute twice
+        for _ in range(2):
+            local_full_params = [
+                torch.zeros(s, device="cuda").fill_(rank)
+                for i, s in enumerate(shapes)
+                if i % 4 == rank
+            ]
+            ctx.redistribute_and_update_params(local_full_params)
+
+        # Verify the same buffer is used (no re-allocation)
+        self.assertEqual(ctx._recv_buffer.data_ptr(), recv_buffer_ptr)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_param_recv_info_completeness(self) -> None:
+        """Verify _param_recv_info has valid entries for all params."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        rank = dist.get_rank()
+        num_params = 5
+        shapes = generate_param_shapes(num_params)
+        params = [torch.zeros(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+        assigned_mask = tuple(i % 4 == rank for i in range(num_params))
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        ctx = RedistributeParamsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_mask,
+            dist_group=dist_group,
+        )
+
+        # Every param should have a valid recv info entry (offset >= 0)
+        for param_idx in range(num_params):
+            offset, chunk_size = ctx._param_recv_info[param_idx]
+            self.assertGreaterEqual(  # type: ignore
+                offset, 0, f"Param {param_idx} has invalid recv offset {offset}"
+            )
+            self.assertGreaterEqual(
+                chunk_size,
+                0,
+                f"Param {param_idx} has invalid recv chunk_size {chunk_size}",
+            )
+
+
+@instantiate_parametrized_tests
+class ComputeParamChunkSizesTest(DTensorTestBase):
+    """Tests for _compute_param_chunk_sizes with actual DTensors."""
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize("num_params", (1, 4, 7))
+    def test_chunk_sizes_match_local_numels(self, num_params: int) -> None:
+        """Verify computed chunk sizes match actual local shard numels."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        shapes = generate_param_shapes(num_params)
+        params = [torch.randn(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+
+        chunk_sizes = _compute_param_chunk_sizes(dtensor_params, self.world_size)
+
+        rank = dist.get_rank()
+        for i, param in enumerate(dtensor_params):
+            local_numel = param.to_local().numel()
+            self.assertEqual(
+                chunk_sizes[i][rank],
+                local_numel,
+                f"Chunk size mismatch for param {i} on rank {rank}",
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_chunk_sizes_sum_to_global_numel(self) -> None:
+        """Verify chunk sizes across all ranks sum to global numel."""
+        device_mesh = init_device_mesh("cuda", (4,))
+        shapes = [(8, 4), (3, 5), (1, 2)]
+        params = [torch.randn(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+
+        chunk_sizes = _compute_param_chunk_sizes(dtensor_params, self.world_size)
+
+        for i, param in enumerate(dtensor_params):
+            global_numel = param.numel()
+            computed_total = sum(chunk_sizes[i])
+            self.assertEqual(
+                computed_total,
+                global_numel,
+                f"Chunk sizes for param {i} sum to {computed_total}, expected {global_numel}",
+            )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Skip when CUDA is not available")
+@instantiate_parametrized_tests
+class GatherGradientsContextTest(DTensorTestBase):
+    """Tests for GatherGradientsContext gradient gathering using all_to_all."""
+
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    def _create_dtensor_params_with_grads(
+        self,
+        shapes: list[tuple[int, ...]],
+        device_mesh: torch.distributed.device_mesh.DeviceMesh,
+    ) -> tuple[torch.distributed.tensor.DTensor, ...]:
+        """Create DTensor params and set their gradients via a dummy backward pass."""
+        dtensor_params = []
+        for shape in shapes:
+            global_tensor = torch.randn(shape, device="cuda", requires_grad=True)
+            dtensor = distribute_tensor(global_tensor, device_mesh, [Shard(0)])
+            dtensor_params.append(dtensor)
+
+        loss = sum(p.sum() for p in dtensor_params)
+        loss.backward()  # type: ignore
+
+        return tuple(dtensor_params)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize("num_params", (1, 4, 7))
+    def test_gather_gradients_matches_full_tensor(self, num_params: int) -> None:
+        """Verify gathered grads match per-param full_tensor() for assigned params."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(num_params)
+        dtensor_params = self._create_dtensor_params_with_grads(shapes, device_mesh)
+
+        # Compute expected full grads using full_tensor()
+        expected_full_grads = [
+            None if p.grad is None else p.grad.full_tensor()  # type: ignore
+            for p in dtensor_params
+        ]
+
+        # Build round-robin assignment mask
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(num_params)
+        )
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+        gathered_grads = ctx.gather_gradients()
+
+        # Verify length
+        self.assertEqual(len(gathered_grads), num_params)
+
+        # Verify correctness for assigned params
+        for i in range(num_params):
+            if i % self.world_size == rank:
+                self.assertIsNotNone(  # type: ignore
+                    gathered_grads[i],
+                    f"Assigned param {i} should have a gathered gradient",
+                )
+                torch.testing.assert_close(
+                    gathered_grads[i],
+                    expected_full_grads[i],
+                    msg=f"Gathered gradient for param {i} does not match full_tensor()",
+                )
+            else:
+                # Unassigned params should be None
+                self.assertIsNone(  # type: ignore
+                    gathered_grads[i],
+                    f"Unassigned param {i} should be None on rank {rank}",
+                )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_gather_gradients_with_none_grads(self) -> None:
+        """Verify params without gradients return None."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(4)
+        # Create params WITHOUT gradients
+        dtensor_params = tuple(
+            distribute_tensor(torch.randn(s, device="cuda"), device_mesh, [Shard(0)])
+            for s in shapes
+        )
+
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(len(shapes))
+        )
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+        gathered_grads = ctx.gather_gradients()
+
+        # All should be None since no gradients were set
+        for i, grad in enumerate(gathered_grads):
+            self.assertIsNone(  # type: ignore
+                grad, f"Param {i} has no grad, should be None"
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_gather_gradients_with_partial_none_grads(self) -> None:
+        """Verify correct handling when some params have grads and some don't."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(4)
+        dtensor_params_list = [
+            distribute_tensor(
+                torch.randn(s, device="cuda", requires_grad=True),
+                device_mesh,
+                [Shard(0)],
+            )
+            for s in shapes
+        ]
+
+        # Only set gradients on even-indexed params (0, 2)
+        for i, p in enumerate(dtensor_params_list):
+            if i % 2 == 0:
+                fake_loss = p.sum()
+                fake_loss.backward()
+
+        dtensor_params = tuple(dtensor_params_list)
+
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(len(shapes))
+        )
+
+        # Compute expected grads via full_tensor() for params that have grads
+        expected_full_grads = [
+            None if p.grad is None else p.grad.full_tensor()  # type: ignore
+            for p in dtensor_params
+        ]
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+        gathered_grads = ctx.gather_gradients()
+
+        for i in range(len(shapes)):
+            if i % self.world_size == rank:
+                if expected_full_grads[i] is not None:
+                    self.assertIsNotNone(gathered_grads[i])  # type: ignore
+                    torch.testing.assert_close(
+                        gathered_grads[i], expected_full_grads[i]
+                    )
+                else:
+                    self.assertIsNone(gathered_grads[i])  # type: ignore
+            else:
+                self.assertIsNone(gathered_grads[i])  # type: ignore
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_gather_gradients_preserves_shape(self) -> None:
+        """Verify gathered gradients have the correct global shape."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = [(8, 4), (12, 3), (5, 7, 2)]
+        # pyrefly: ignore [bad-argument-type]
+        dtensor_params = self._create_dtensor_params_with_grads(shapes, device_mesh)
+
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(len(shapes))
+        )
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+        gathered_grads = ctx.gather_gradients()
+
+        for i in range(len(shapes)):
+            if i % self.world_size == rank:
+                self.assertIsNotNone(gathered_grads[i])  # type: ignore
+                self.assertEqual(
+                    gathered_grads[i].shape,  # type: ignore
+                    torch.Size(shapes[i]),
+                    f"Shape mismatch for param {i}: expected {shapes[i]}, got {gathered_grads[i].shape}",  # type: ignore
+                )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_gather_gradients_multiple_calls(self) -> None:
+        """Verify gather_gradients works correctly across multiple calls."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(4)
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(len(shapes))
+        )
+
+        # First call with initial gradients
+        dtensor_params = self._create_dtensor_params_with_grads(shapes, device_mesh)
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+
+        first_grads = ctx.gather_gradients()
+
+        # Zero gradients and do another backward with different values
+        for p in dtensor_params:
+            if p.grad is not None:
+                p.grad.zero_()
+
+        loss = sum(2.0 * p.sum() for p in dtensor_params)
+        loss.backward()  # type: ignore
+
+        expected_second_grads = [
+            None if p.grad is None else p.grad.full_tensor()  # type: ignore
+            for p in dtensor_params
+        ]
+
+        second_grads = ctx.gather_gradients()
+
+        # Verify second call produces correct results (different from first)
+        for i in range(len(shapes)):
+            if i % self.world_size == rank:
+                self.assertIsNotNone(second_grads[i])  # type: ignore
+                torch.testing.assert_close(second_grads[i], expected_second_grads[i])
+                # Verify second call differs from first
+                # (2x gradient vs 1x gradient for sum)
+                self.assertFalse(  # type: ignore
+                    torch.equal(first_grads[i], second_grads[i]),  # type: ignore
+                    f"Second gather should differ from first for param {i}",
+                )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_empty_params_raises_error(self) -> None:
+        """Verify empty params raises ValueError."""
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        with self.assertRaises(ValueError):
+            GatherGradientsContext(
+                params=(),
+                assigned_params_mask=(),
+                dist_group=dist_group,
+            )
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    def test_metadata_precomputation(self) -> None:
+        """Verify precomputed metadata is consistent."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(4)
+        dtensor_params = tuple(
+            distribute_tensor(torch.randn(s, device="cuda"), device_mesh, [Shard(0)])
+            for s in shapes
+        )
+
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(len(shapes))
+        )
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+
+        # Verify send/recv sizes have correct count
+        self.assertEqual(len(ctx._send_sizes), self.world_size)
+        self.assertEqual(len(ctx._recv_sizes), self.world_size)
+
+        # Total send size should equal sum of this rank's chunk sizes for all params
+        expected_total_send = sum(
+            ctx._param_chunk_sizes[i][rank] for i in range(len(shapes))
+        )
+        self.assertEqual(ctx._total_send_size, expected_total_send)
+
+        # Verify recv buffer is pre-allocated with correct size
+        self.assertEqual(ctx._recv_buffer.numel(), ctx._total_recv_size)
+
+        # Verify local param indices match assignment mask
+        expected_local_indices = [
+            i for i, assigned in enumerate(assigned_params_mask) if assigned
+        ]
+        self.assertEqual(ctx._local_param_indices, expected_local_indices)
+
+    @with_comms
+    @skip_if_lt_x_gpu(4)
+    @parametrize("num_params", (1, 4, 7))
+    def test_gather_params_matches_full_tensor(self, num_params: int) -> None:
+        """Verify gather_params() returns correct full parameter tensors."""
+        device_mesh = init_device_mesh("cuda", (self.world_size,))
+        rank = dist.get_rank()
+        dist_group = dist.distributed_c10d._get_default_group()
+
+        shapes = generate_param_shapes(num_params)
+        # Create params with distinct values so we can verify correctness
+        params = [torch.randn(s, device="cuda") for s in shapes]
+        dtensor_params = tuple(
+            distribute_tensor(t, device_mesh, [Shard(0)]) for t in params
+        )
+
+        # Compute expected full params using full_tensor()
+        expected_full_params = [p.full_tensor() for p in dtensor_params]
+
+        assigned_params_mask = tuple(
+            i % self.world_size == rank for i in range(num_params)
+        )
+
+        ctx = GatherGradientsContext(
+            params=dtensor_params,
+            assigned_params_mask=assigned_params_mask,
+            dist_group=dist_group,
+        )
+        gathered_params = ctx.gather_params()
+
+        # Verify length
+        self.assertEqual(len(gathered_params), num_params)
+
+        # Verify correctness for assigned params
+        for i in range(num_params):
+            if i % self.world_size == rank:
+                self.assertIsNotNone(  # type: ignore
+                    gathered_params[i],
+                    f"Assigned param {i} should have a gathered value",
+                )
+                torch.testing.assert_close(
+                    gathered_params[i],
+                    expected_full_params[i],
+                    msg=f"Gathered param {i} does not match full_tensor()",
+                )
+            else:
+                self.assertIsNone(  # type: ignore
+                    gathered_params[i],
+                    f"Unassigned param {i} should be None on rank {rank}",
+                )
diff --git a/distributed_shampoo/utils/load_balancing_utils.py b/distributed_shampoo/utils/load_balancing_utils.py
index a052b8a..684f803 100644
--- a/distributed_shampoo/utils/load_balancing_utils.py
+++ b/distributed_shampoo/utils/load_balancing_utils.py
@@ -54,7 +54,7 @@ class PolynomialComputationalCostModel(CostModel):
 
     def cost(self, tensor: torch.Tensor) -> float:
         return sum(
-            max(self.min_cost, polyval(x=dim_size, c=self.coefficients))  # type: ignore[misc,call-overload]
+            max(self.min_cost, float(polyval(x=dim_size, c=self.coefficients)))
             for dim_size in tensor.shape
         )
 
diff --git a/distributed_shampoo/utils/optimizer_modules.py b/distributed_shampoo/utils/optimizer_modules.py
index c9506d6..5de1ee9 100644
--- a/distributed_shampoo/utils/optimizer_modules.py
+++ b/distributed_shampoo/utils/optimizer_modules.py
@@ -119,7 +119,7 @@ def save_to_state_dict(
         save_to_state_dict(self.__dict__.items(), destination)
         return destination
 
-    def load_state_dict(
+    def load_state_dict(  # noqa: C901
         self, state_dict: StateDict, store_non_tensors: bool = False
     ) -> None:
         """
@@ -214,12 +214,15 @@ def load_from_new_state_to_old_state(
                 # When state dict is flattened/unflattened, dictionary keys in new_state become strings
                 if isinstance(new_state, dict) and len(new_state) > 0:
                     # Convert new_state dict to match the collection structure expected by old_state
+                    # pyrefly: ignore [bad-assignment]
                     new_state = _convert_state_key_from_str_to_int(old_state, new_state)
 
                 old_state = type(old_state)(
                     (
                         load_from_new_state_to_old_state(
+                            # pyrefly: ignore [bad-argument-type]
                             old_state=old_value,
+                            # pyrefly: ignore [bad-index]
                             new_state=new_state[i],
                         )
                         if store_non_tensors
diff --git a/distributed_shampoo/utils/shampoo_fully_shard_utils.py b/distributed_shampoo/utils/shampoo_fully_shard_utils.py
new file mode 100644
index 0000000..789a194
--- /dev/null
+++ b/distributed_shampoo/utils/shampoo_fully_shard_utils.py
@@ -0,0 +1,833 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
+
+import logging
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from itertools import accumulate
+
+import torch
+from torch import distributed as dist, Tensor
+from torch.distributed.tensor import DTensor
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def prepare_update_param_buffers(
+    params: tuple[DTensor, ...], group_size: int
+) -> list[Tensor]:
+    """Allocates a persistent shadow copy of updated parameters."""
+    if any(p.dtype != params[0].dtype for p in params):
+        raise NotImplementedError(
+            "When using round-robin assignment in FSDP Shampoo, parameters of "
+            "different dtypes are not currently supported."
+        )
+
+    param_sizes = [p.to_local().numel() for p in params]
+    buffer_size = sum(param_sizes)
+    buffer = params[0].to_local().new_zeros(buffer_size)
+    buffer_offsets = list(accumulate(param_sizes))
+
+    def round_up_to_multiple_of(x: int, y: int) -> int:
+        return ((x + y - 1) // y) * y
+
+    pad_len = round_up_to_multiple_of(len(buffer_offsets), group_size) - len(
+        buffer_offsets
+    )
+
+    # The padding logic below handles when a rank has some parameters but fewer than group size.
+    # For example, if group size is 4 and there are 3 parameters, it will pad a 0-sized tensor at the end.
+    # Example:
+    #   Assume we have 3 parameters and group size is 4. param0: 100, param1: 200, param2: 300.
+    #   buffer_offsets = [100, 300, 600, 600] (note that the last element is 600)
+    #   This buffer for communication have 4 chunks.
+    #   - Rank 0: [0, 100)
+    #   - Rank 1: [100, 300)
+    #   - Rank 2: [300, 600)
+    #   - Rank 3: [600, 600) (empty tensor)
+    # Pad the list with empty tensors to ensure each rank participates in all-to-all.
+    buffer_offsets.extend([buffer_size] * pad_len)
+    # Drop the last element as torch.tensor_split takes indices as split points.
+    buffer_offsets = buffer_offsets[:-1]
+
+    return list(torch.tensor_split(buffer, buffer_offsets))
+
+
+def redistribute_and_update_params(
+    params: tuple[DTensor, ...],
+    local_full_params: list[Tensor],
+    update_param_buffers: list[Tensor],
+    dist_group: torch.distributed.ProcessGroup,
+) -> None:
+    """Redistributes updated parameters to each parameter's rank."""
+    group_size = dist_group.size()
+
+    # Run all-to-all collectives to exchange the updated parameters across
+    # ranks in group. This implementation runs multiple rounds of a2a ops
+    # if the number of parameters is larger than the world size.
+    a2a_rounds = range(len(update_param_buffers) // group_size)
+    logger.info(f"Running {len(a2a_rounds)} rounds of a2a ops.")
+    for a2a_round in a2a_rounds:
+        # Send either a valid full parameter, or a padding zero tensor.
+        send_param = (
+            local_full_params[a2a_round]
+            if a2a_round < len(local_full_params)
+            else params[0].to_local().new_zeros(0)
+        )
+        # Chunk the send_param to exactly group_size slices to distribute to
+        # all ranks. We need to manually pad the result of torch.chunk since
+        # it does not guarantee that the result has the desired chunks.
+        send_list = [t.flatten() for t in torch.chunk(send_param, group_size, dim=0)]
+        if len(send_list) < group_size:
+            # NOTE: Intentionally use `torch.tensor_split` here to do a trivial
+            # split to ensure that the padding is in contiguous memory space as
+            # is required for all-to-all collectives.
+            append_len = group_size - len(send_list)
+            last_t = send_list[-1]
+            split_indices = [send_list[-1].shape[0]] * append_len
+            send_list.extend(torch.tensor_split(last_t, split_indices, dim=0)[1:])
+        assert len(send_list) == group_size
+
+        # Specify receive list as a range of update_param_buffers.
+        recv_list = update_param_buffers[
+            a2a_round * group_size : (a2a_round + 1) * group_size
+        ]
+
+        dist.all_to_all(recv_list, send_list, dist_group)
+
+    torch._foreach_copy_(
+        [p.to_local().flatten() for p in params], update_param_buffers[: len(params)]
+    )
+
+
+def _compute_chunk_sizes(numel: int, num_chunks: int) -> list[int]:
+    """Compute chunk sizes that torch.chunk would produce.
+
+    torch.chunk(tensor, n, dim) semantics for n chunks along dim with size d:
+    - Computes chunk_size = ceil(d / n)
+    - First (n-1) chunks get chunk_size elements, last chunk gets remainder
+    - If d < n, only d chunks are produced (each of size 1)
+
+    Args:
+        numel (int): Size of the dimension to split.
+        num_chunks (int): Number of chunks to split into.
+
+    Returns:
+        sizes (list[int]): List of chunk sizes, one per chunk.
+    """
+    if numel == 0:
+        return [0] * num_chunks
+
+    if numel >= num_chunks:
+        # torch.chunk computes ceil(numel / num_chunks) as the chunk size
+        chunk_size = (numel + num_chunks - 1) // num_chunks  # This is ceil division
+        # First chunks get chunk_size, last chunk gets remainder
+        sizes = []
+        remaining = numel
+        for _i in range(num_chunks):
+            if remaining <= 0:
+                sizes.append(0)
+            elif remaining >= chunk_size:
+                sizes.append(chunk_size)
+                remaining -= chunk_size
+            else:
+                sizes.append(remaining)
+                remaining = 0
+        return sizes
+    else:
+        # numel < num_chunks: each element becomes its own chunk
+        # Ranks beyond numel get empty chunks (size 0)
+        sizes = [1] * numel + [0] * (num_chunks - numel)
+        return sizes
+
+
+def _compute_param_chunk_sizes(
+    params: tuple[DTensor, ...], group_size: int
+) -> list[list[int]]:
+    """Compute per-rank chunk sizes for all params based on FSDP dim-0 sharding.
+
+    For each parameter, computes how many elements each rank holds after
+    FSDP's dim-0 sharding (torch.chunk semantics).
+
+    Args:
+        params (tuple[DTensor, ...]): Tuple of DTensor parameters.
+        group_size (int): Number of ranks in the process group.
+
+    Returns:
+        param_chunk_sizes (list[list[int]]): List of lists, where param_chunk_sizes[i][r] is the
+            number of elements rank r holds for parameter i.
+    """
+    param_chunk_sizes: list[list[int]] = []
+    for param in params:
+        global_shape = tuple(param.shape)
+
+        if len(global_shape) == 0:
+            # Scalar tensor - all on rank 0
+            chunk_sizes = [1] + [0] * (group_size - 1)
+        else:
+            dim0_size = global_shape[0]
+            remaining_numel = math.prod(global_shape[1:])
+
+            # Compute dim-0 chunk sizes using torch.chunk semantics
+            dim0_chunks = _compute_chunk_sizes(dim0_size, group_size)
+            # Multiply by remaining dimensions to get actual element counts
+            chunk_sizes = [c * remaining_numel for c in dim0_chunks]
+
+        param_chunk_sizes.append(chunk_sizes)
+
+    return param_chunk_sizes
+
+
+def _build_buffer_views(buffer: Tensor, sizes: list[int]) -> list[Tensor]:
+    """Build contiguous views into a flat buffer given per-section sizes."""
+    views: list[Tensor] = []
+    offset = 0
+    for size in sizes:
+        views.append(buffer[offset : offset + size])
+        offset += size
+    return views
+
+
+# NOTE: Plans group parallel per-entry lists; the lists themselves stay flat
+# so the hot path passes them directly to torch._foreach_copy_.
+
+
+@dataclass
+class _RedistributeSendPlan:
+    """Per-entry send metadata for RedistributeParamsContext.
+
+    Source data comes from the caller (local_full_params); only destination
+    views and routing metadata are stored here.
+    """
+
+    param_indices: list[int] = field(default_factory=list)
+    peer_ranks: list[int] = field(default_factory=list)
+    dst_views: list[Tensor] = field(default_factory=list)
+
+
+@dataclass
+class _RedistributeRecvPlan:
+    """Per-entry recv metadata for RedistributeParamsContext.
+
+    Source views point into the recv buffer; destination views point into the
+    local param shards. Iteration is src -> dst.
+    """
+
+    src_views: list[Tensor] = field(default_factory=list)
+    dst_views: list[Tensor] = field(default_factory=list)
+
+
+@dataclass
+class _GatherSendPlan:
+    """Per-entry send metadata for GatherGradientsContext.
+
+    Source data comes from a callable (gradient or param value); only the
+    routing metadata (param indices, chunk sizes) and destination send-buffer
+    views are stored here.
+    """
+
+    param_indices: list[int] = field(default_factory=list)
+    chunk_sizes: list[int] = field(default_factory=list)
+    dst_views: list[Tensor] = field(default_factory=list)
+
+
+@dataclass
+class _GatherUnpackPlan:
+    """Per-entry recv-buffer-unpack metadata for GatherGradientsContext.
+
+    Source views point into the recv buffer; destination views point into the
+    pre-allocated full-grad buffers. Iteration is src -> dst.
+    """
+
+    src_views: list[Tensor] = field(default_factory=list)
+    dst_views: list[Tensor] = field(default_factory=list)
+
+
+class AllToAllContext(ABC):
+    """Base class for coalesced all_to_all communication contexts.
+
+    Provides shared infrastructure for precomputing FSDP dim-0 sharding metadata,
+    allocating persistent send/recv buffers, and executing a single all_to_all
+    collective. Subclasses implement the direction-specific send/recv size
+    computation and additional metadata precomputation.
+
+    TODO(irisz): Currently assumes dim-0 sharding only. FSDP2 supports non-dim-0 sharding
+    but it is not enabled on APS. Extend to support arbitrary shard dimensions if needed.
+
+    Args:
+        params (tuple[DTensor, ...]): Tuple of DTensor parameters.
+        assigned_params_mask (tuple[bool, ...]): Boolean mask indicating which params this rank owns.
+        dist_group (torch.distributed.ProcessGroup): Process group for communication.
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        params: tuple[DTensor, ...],
+        assigned_params_mask: tuple[bool, ...],
+        dist_group: torch.distributed.ProcessGroup,
+    ) -> None:
+        self._params: tuple[DTensor, ...] = params
+        self._assigned_params_mask: tuple[bool, ...] = assigned_params_mask
+        self._dist_group: torch.distributed.ProcessGroup = dist_group
+        self._group_size: int = dist_group.size()
+        self._rank: int = dist.get_rank(group=dist_group)
+
+        if not params:
+            raise ValueError("params cannot be empty")
+
+        if len(params) != len(assigned_params_mask):
+            raise ValueError(
+                f"len(params) ({len(params)}) != "
+                f"len(assigned_params_mask) ({len(assigned_params_mask)})"
+            )
+
+        self._dtype: torch.dtype = params[0].to_local().dtype
+        self._device: torch.device = params[0].to_local().device
+
+        # Validate all params have same dtype
+        if any(p.to_local().dtype != self._dtype for p in params):
+            raise NotImplementedError(
+                "Parameters of different dtypes are not currently supported."
+            )
+
+        self._precompute_metadata()
+
+    def _precompute_metadata(self) -> None:
+        """Precompute shared metadata for all_to_all communication."""
+        group_size = self._group_size
+
+        # Compute which params this rank owns (based on round-robin assignment)
+        self._local_param_indices: list[int] = [
+            i for i, assigned in enumerate(self._assigned_params_mask) if assigned
+        ]
+
+        # Compute chunk sizes using FSDP's dim-0 sharding semantics
+        self._param_chunk_sizes: list[list[int]] = _compute_param_chunk_sizes(
+            self._params, group_size
+        )
+
+        # Compute direction-specific send/recv sizes
+        self._send_sizes, self._recv_sizes = self._compute_send_recv_sizes()
+        self._total_send_size = sum(self._send_sizes)
+        self._total_recv_size = sum(self._recv_sizes)
+
+        # Pre-allocate persistent send and recv buffers for all_to_all.
+        self._send_buffer = torch.empty(
+            self._total_send_size, dtype=self._dtype, device=self._device
+        )
+        self._recv_buffer = torch.empty(
+            self._total_recv_size, dtype=self._dtype, device=self._device
+        )
+
+        # Pre-compute send_list and recv_list as contiguous views into buffers.
+        self._send_list = _build_buffer_views(self._send_buffer, self._send_sizes)
+        self._recv_list = _build_buffer_views(self._recv_buffer, self._recv_sizes)
+
+        # Delegate subclass-specific precomputation.
+        self._precompute_subclass_metadata()
+
+    @abstractmethod
+    def _compute_send_recv_sizes(self) -> tuple[list[int], list[int]]:
+        """Compute per-peer send and recv sizes. Subclasses swap the direction."""
+        ...
+
+    @abstractmethod
+    def _precompute_subclass_metadata(self) -> None:
+        """Hook for subclass-specific metadata precomputation."""
+        ...
+
+    def _execute_all_to_all(self) -> None:
+        """Execute all_to_all using persistent send/recv buffers."""
+        dist.all_to_all(self._recv_list, self._send_list, group=self._dist_group)
+
+
+class RedistributeParamsContext(AllToAllContext):
+    """Context for optimized parameter redistribution using all_to_all.
+
+    Sends full params from owning ranks to all ranks' local shards.
+    Precomputes all static information at init, then uses a single all_to_all
+    collective per step instead of multiple point-to-point calls.
+
+    Example:
+        # During initialization:
+        ctx = RedistributeParamsContext(params, assigned_mask, dist_group)
+
+        # During each optimizer step:
+        ctx.redistribute_and_update_params(local_full_params)
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        params: tuple[DTensor, ...],
+        assigned_params_mask: tuple[bool, ...],
+        dist_group: torch.distributed.ProcessGroup,
+    ) -> None:
+        # Forward-declare subclass-specific attributes for Pyre before
+        # super().__init__() calls _precompute_subclass_metadata().
+        self._param_local_numels: list[int] = []
+        self._param_global_numels: list[int] = []
+        self._param_to_local_idx: dict[int, int] = {}
+        self._local_param_global_shapes: list[tuple[int, ...]] = []
+        self._param_recv_info: list[tuple[int, int]] = []
+        self._send_plan = _RedistributeSendPlan()
+        self._recv_plan = _RedistributeRecvPlan()
+        super().__init__(params, assigned_params_mask, dist_group)
+
+    def _compute_send_recv_sizes(self) -> tuple[list[int], list[int]]:
+        """Owner-to-all: send owned params' chunks to each peer, recv each peer's chunks."""
+        group_size = self._group_size
+
+        # Send sizes: sum of chunk sizes for owned params going to each peer
+        send_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            send_size = sum(
+                self._param_chunk_sizes[param_idx][peer_rank]
+                for param_idx in self._local_param_indices
+            )
+            send_sizes.append(send_size)
+
+        # Recv sizes: sum of chunk sizes for params owned by each peer
+        recv_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            recv_size = sum(
+                self._param_chunk_sizes[param_idx][self._rank]
+                for param_idx in peer_param_indices
+            )
+            recv_sizes.append(recv_size)
+
+        return send_sizes, recv_sizes
+
+    def _precompute_subclass_metadata(self) -> None:
+        """Precompute redistribute-specific views and metadata."""
+        self._param_local_numels = [p.to_local().numel() for p in self._params]
+        self._param_global_numels = [p.numel() for p in self._params]
+        self._compute_send_views()
+        self._compute_recv_views()
+
+        logger.info(
+            f"RedistributeParamsContext initialized: "
+            f"rank={self._rank}, group_size={self._group_size}, "
+            f"num_params={len(self._params)}, "
+            f"local_params={len(self._local_param_indices)}, "
+            f"send_sizes={self._send_sizes}, "
+            f"recv_sizes={self._recv_sizes}, "
+            f"param_local_numels={self._param_local_numels}, "
+            f"param_global_numels={self._param_global_numels}, "
+            f"param_chunk_sizes={self._param_chunk_sizes}"
+        )
+
+    def _compute_send_views(self) -> None:
+        """Pre-compute per-(local_param, peer_rank) destination views into send buffer."""
+        group_size = self._group_size
+        send_offset = 0
+        for peer_rank in range(group_size):
+            for param_idx in self._local_param_indices:
+                chunk_size = self._param_chunk_sizes[param_idx][peer_rank]
+                if chunk_size > 0:
+                    self._send_plan.param_indices.append(param_idx)
+                    self._send_plan.peer_ranks.append(peer_rank)
+                    self._send_plan.dst_views.append(
+                        self._send_buffer[send_offset : send_offset + chunk_size]
+                    )
+                    send_offset += chunk_size
+
+        # Build a mapping from param_idx to local_idx for fast lookup.
+        self._param_to_local_idx.update(
+            {
+                param_idx: local_idx
+                for local_idx, param_idx in enumerate(self._local_param_indices)
+            }
+        )
+
+        # Pre-compute global shapes for owned params (used for chunking).
+        self._local_param_global_shapes.extend(
+            tuple(self._params[param_idx].shape)
+            for param_idx in self._local_param_indices
+        )
+
+    def _compute_recv_views(self) -> None:
+        """Pre-compute recv unpacking info and copy-out views."""
+        group_size = self._group_size
+
+        # Precompute recv unpacking info
+        # Recv buffer layout: [data from rank 0][data from rank 1]...
+        # _param_recv_info[param_idx] = (offset in recv_buffer, chunk_size)
+        self._param_recv_info = [(-1, -1)] * len(self._params)
+        recv_offset = 0
+        for peer_rank in range(group_size):
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            for param_idx in peer_param_indices:
+                chunk_size = self._param_chunk_sizes[param_idx][self._rank]
+                self._param_recv_info[param_idx] = (recv_offset, chunk_size)
+                recv_offset += chunk_size
+
+        # Pre-compute recv copy-out views for batched copy in
+        # redistribute_and_update_params().
+        for param_idx in range(len(self._params)):
+            recv_offset, chunk_size = self._param_recv_info[param_idx]
+            if chunk_size > 0:
+                local_param = self._params[param_idx].to_local()
+                local_numel = local_param.numel()
+                assert chunk_size == local_numel, (
+                    f"chunk_size ({chunk_size}) != local_numel ({local_numel}) "
+                    f"for param {param_idx}"
+                )
+                self._recv_plan.src_views.append(
+                    self._recv_buffer[recv_offset : recv_offset + chunk_size]
+                )
+                self._recv_plan.dst_views.append(local_param.flatten())
+
+    @torch.no_grad()
+    def redistribute_and_update_params(
+        self,
+        local_full_params: list[Tensor],
+    ) -> None:
+        """Redistribute updated parameters using a SINGLE coalesced all_to_all.
+
+        This combines ALL params into one collective call instead of multiple rounds,
+        reducing the number of collective operations.
+
+        Args:
+            local_full_params (list[Tensor]): List of full updated params computed by this rank.
+                Must match the params indicated by assigned_params_mask.
+        """
+        assert len(local_full_params) == len(self._local_param_indices), (
+            f"Expected {len(self._local_param_indices)} local params, "
+            f"got {len(local_full_params)}"
+        )
+
+        group_size = self._group_size
+
+        # Pre-chunk all local full params by peer rank for source view construction.
+        # param_chunks[local_idx][peer_rank] = flattened chunk tensor
+        param_chunks: list[list[Tensor]] = []
+        for local_idx in range(len(self._local_param_indices)):
+            full_param = local_full_params[local_idx]
+            global_shape = self._local_param_global_shapes[local_idx]
+
+            if len(global_shape) == 0:
+                # Scalar tensor - all goes to rank 0.
+                # Note: FSDP2 does not produce scalar params in practice;
+                # this branch is defensive.
+                chunks = [full_param.flatten()]
+                while len(chunks) < group_size:
+                    chunks.append(
+                        torch.empty(0, dtype=full_param.dtype, device=full_param.device)
+                    )
+            else:
+                full_param_reshaped = full_param.view(global_shape)
+                dim0_chunks = torch.chunk(full_param_reshaped, group_size, dim=0)
+                chunks = [c.flatten() for c in dim0_chunks]
+                while len(chunks) < group_size:
+                    chunks.append(
+                        torch.empty(0, dtype=full_param.dtype, device=full_param.device)
+                    )
+            param_chunks.append(chunks)
+
+        # Build source list matching the pre-computed destination views and
+        # batch-copy into the send buffer.
+        src_list: list[Tensor] = []
+        dst_list: list[Tensor] = []
+        for param_idx, peer_rank, dst_view in zip(
+            self._send_plan.param_indices,
+            self._send_plan.peer_ranks,
+            self._send_plan.dst_views,
+            strict=True,
+        ):
+            local_idx = self._param_to_local_idx[param_idx]
+            src_list.append(param_chunks[local_idx][peer_rank])
+            dst_list.append(dst_view)
+
+        if dst_list:
+            torch._foreach_copy_(dst_list, src_list)
+
+        # Execute SINGLE all_to_all - this is the key optimization!
+        dist.all_to_all(self._recv_list, self._send_list, group=self._dist_group)
+
+        # Batch-copy received chunks to local param shards in one fused kernel.
+        if self._recv_plan.dst_views:
+            torch._foreach_copy_(self._recv_plan.dst_views, self._recv_plan.src_views)
+
+
+class GatherGradientsContext(AllToAllContext):
+    """Context for gathering gradients to owning ranks using all_to_all.
+
+    This is the mirror of RedistributeParamsContext. While RedistributeParamsContext
+    sends full params from owning ranks to all ranks' local shards,
+    GatherGradientsContext gathers local gradient shards from all ranks to the
+    owning rank to reconstruct full gradients.
+
+    Data flow:
+        Each rank has local grad shards for ALL params (from FSDP).
+        Each rank sends its local shard of param i to the rank that owns param i.
+        The owning rank concatenates received shards along dim-0 to get the full gradient.
+
+    This replaces per-param full_tensor() all-gathers with a single all_to_all,
+    reducing peak memory from O(total_params) to O(total_params / world_size).
+
+    Example:
+        # During initialization:
+        ctx = GatherGradientsContext(params, assigned_mask, dist_group)
+
+        # During each optimizer step:
+        full_grads = ctx.gather_gradients()
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        params: tuple[DTensor, ...],
+        assigned_params_mask: tuple[bool, ...],
+        dist_group: torch.distributed.ProcessGroup,
+    ) -> None:
+        # Forward-declare subclass-specific attributes for Pyre before
+        # super().__init__() calls _precompute_subclass_metadata().
+        self._param_global_shapes: list[tuple[int, ...]] = []
+        self._param_recv_info: list[list[tuple[int, int]]] = []
+        self._param_local_idx_map: dict[int, int] = {}
+        self._full_grad_buffers: list[Tensor] = []
+        self._send_plan = _GatherSendPlan()
+        self._unpack_plan = _GatherUnpackPlan()
+        super().__init__(params, assigned_params_mask, dist_group)
+
+    def _compute_send_recv_sizes(self) -> tuple[list[int], list[int]]:
+        """All-to-owner: send this rank's chunks to each owner, recv all peers' chunks for owned params."""
+        group_size = self._group_size
+
+        # Send sizes: for each peer rank, sum of this rank's local shard sizes
+        # for params owned by that peer.
+        send_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            send_size = sum(
+                self._param_chunk_sizes[param_idx][self._rank]
+                for param_idx in peer_param_indices
+            )
+            send_sizes.append(send_size)
+
+        # Recv sizes: for each peer rank, sum of that peer's shard sizes
+        # for params owned by this rank.
+        recv_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            recv_size = sum(
+                self._param_chunk_sizes[param_idx][peer_rank]
+                for param_idx in self._local_param_indices
+            )
+            recv_sizes.append(recv_size)
+
+        return send_sizes, recv_sizes
+
+    def _precompute_subclass_metadata(self) -> None:
+        """Precompute gather-specific views and metadata."""
+        group_size = self._group_size
+
+        # Store global shapes for all params (used to reconstruct full grads)
+        self._param_global_shapes = [tuple(p.shape) for p in self._params]
+
+        # Precompute recv unpacking info for each assigned param.
+        # _param_recv_info[local_idx] = list of (offset, size) for each peer rank
+        recv_section_starts: list[int] = []
+        offset = 0
+        for recv_size in self._recv_sizes:
+            recv_section_starts.append(offset)
+            offset += recv_size
+
+        peer_offsets = list(recv_section_starts)  # mutable copy
+        self._param_local_idx_map = {
+            param_idx: local_idx
+            for local_idx, param_idx in enumerate(self._local_param_indices)
+        }
+        for _ in self._local_param_indices:
+            self._param_recv_info.append([(-1, -1)] * group_size)
+
+        for peer_rank in range(group_size):
+            for param_idx in self._local_param_indices:
+                local_idx = self._param_local_idx_map[param_idx]
+                chunk_size = self._param_chunk_sizes[param_idx][peer_rank]
+                self._param_recv_info[local_idx][peer_rank] = (
+                    peer_offsets[peer_rank],
+                    chunk_size,
+                )
+                peer_offsets[peer_rank] += chunk_size
+
+        self._compute_unpack_views()
+        self._compute_send_dst_views()
+
+        logger.info(
+            f"GatherGradientsContext initialized: "
+            f"rank={self._rank}, group_size={group_size}, "
+            f"num_params={len(self._params)}, "
+            f"local_params={len(self._local_param_indices)}, "
+            f"send_sizes={self._send_sizes}, "
+            f"recv_sizes={self._recv_sizes}"
+        )
+
+    def _compute_unpack_views(self) -> None:
+        """Pre-allocate full gradient buffers and pre-compute recv unpack views."""
+        for local_idx in range(len(self._local_param_indices)):
+            total_size = sum(
+                size for _, size in self._param_recv_info[local_idx] if size > 0
+            )
+            grad_buffer = torch.empty(
+                total_size, dtype=self._dtype, device=self._device
+            )
+            self._full_grad_buffers.append(grad_buffer)
+            dst_offset = 0
+            for recv_offset, size in self._param_recv_info[local_idx]:
+                if size > 0:
+                    self._unpack_plan.src_views.append(
+                        self._recv_buffer[recv_offset : recv_offset + size]
+                    )
+                    self._unpack_plan.dst_views.append(
+                        grad_buffer[dst_offset : dst_offset + size]
+                    )
+                    dst_offset += size
+
+    def _compute_send_dst_views(self) -> None:
+        """Pre-compute destination views into the send buffer for batched copy."""
+        group_size = self._group_size
+        send_offset = 0
+        for peer_rank in range(group_size):
+            for param_idx in range(len(self._params)):
+                if param_idx % group_size != peer_rank:
+                    continue
+                chunk_size = self._param_chunk_sizes[param_idx][self._rank]
+                if chunk_size > 0:
+                    self._send_plan.param_indices.append(param_idx)
+                    self._send_plan.chunk_sizes.append(chunk_size)
+                    self._send_plan.dst_views.append(
+                        self._send_buffer[send_offset : send_offset + chunk_size]
+                    )
+                    send_offset += chunk_size
+
+    def _pack_send_buffer(self, data_source: Callable[[int], Tensor | None]) -> None:
+        """Pack data into the persistent send buffer using foreach_copy_.
+
+        Args:
+            data_source (Callable[[int], Tensor | None]): Callable that takes a param index
+                and returns the local tensor to send (flattened and sliced to chunk_size),
+                or None to zero the destination region.
+        """
+        src_list: list[Tensor] = []
+        dst_list: list[Tensor] = []
+
+        for param_idx, chunk_size, dst in zip(
+            self._send_plan.param_indices,
+            self._send_plan.chunk_sizes,
+            self._send_plan.dst_views,
+            strict=True,
+        ):
+            data = data_source(param_idx)
+            if data is not None:
+                src_list.append(data[:chunk_size])
+                dst_list.append(dst)
+            else:
+                dst.zero_()
+
+        if dst_list:
+            torch._foreach_copy_(dst_list, src_list)
+
+    def _unpack_recv_buffer(self, has_grad: list[bool]) -> list[Tensor | None]:
+        """Unpack the recv buffer into full gradients for assigned params.
+
+        Uses pre-computed views and foreach_copy_ to batch-copy all shards from
+        the recv buffer into pre-allocated full gradient buffers, then reshapes
+        each to the original global shape.
+
+        Args:
+            has_grad (list[bool]): Boolean mask indicating which params have gradients.
+
+        Returns:
+            full_grads (list[Tensor | None]): List of full gradients (same length as self._params).
+                None for params without gradients or unassigned params.
+        """
+        # Batch-copy all recv shards into pre-allocated full grad buffers.
+        if self._unpack_plan.dst_views:
+            torch._foreach_copy_(
+                self._unpack_plan.dst_views, self._unpack_plan.src_views
+            )
+
+        full_grads: list[Tensor | None] = [None for _ in self._params]
+        for local_idx, param_idx in enumerate(self._local_param_indices):
+            if not has_grad[param_idx]:
+                continue
+            global_shape = self._param_global_shapes[param_idx]
+            if len(global_shape) == 0:
+                full_grads[param_idx] = self._full_grad_buffers[local_idx].squeeze()
+            else:
+                full_grads[param_idx] = self._full_grad_buffers[local_idx].view(
+                    global_shape
+                )
+
+        return full_grads
+
+    @torch.no_grad()
+    def gather_gradients(self) -> list[Tensor | None]:
+        """Gather gradients from all ranks using a single all_to_all.
+
+        Each rank sends its local grad shard of each param to the rank that owns
+        that param. The owning rank then concatenates the received shards along
+        dim-0 to reconstruct the full gradient.
+
+        Returns:
+            full_grads: List of full gradients for ALL params (same length as
+                self._params). Entries are None for params not assigned to this
+                rank, regardless of whether they have gradients. Only params
+                assigned to this rank will have non-None values.
+        """
+        has_grad: list[bool] = [p.grad is not None for p in self._params]
+
+        def _grad_source(param_idx: int) -> Tensor | None:
+            if not has_grad[param_idx]:
+                return None
+            local_grad = self._params[param_idx].grad
+            assert local_grad is not None
+            return local_grad.to_local().flatten()  # type: ignore
+
+        self._pack_send_buffer(_grad_source)
+        self._execute_all_to_all()
+
+        result = self._unpack_recv_buffer(has_grad)
+        # Clone to detach from _full_grad_buffers, which will be overwritten
+        # by subsequent gather_gradients() calls.
+        return [t.clone() if t is not None else None for t in result]
+
+    @torch.no_grad()
+    def gather_params(self) -> list[Tensor | None]:
+        """Gather full parameter values from all ranks using a single all_to_all.
+
+        Similar to gather_gradients(), but operates on parameter values instead
+        of gradients. Every parameter always has a value, so no None handling
+        is needed on the send side.
+
+        Returns:
+            full_params: List of full parameter tensors (same length as
+                self._params). Entries are None for params not assigned to this
+                rank. Only params assigned to this rank will have non-None values.
+        """
+        self._pack_send_buffer(
+            lambda param_idx: self._params[param_idx].to_local().flatten()
+        )
+        self._execute_all_to_all()
+
+        result = self._unpack_recv_buffer(has_grad=[True] * len(self._params))
+        # Clone to detach from _full_grad_buffers, which will be overwritten
+        # by subsequent gather_gradients() calls.
+        return [t.clone() if t is not None else None for t in result]
diff --git a/distributed_shampoo/utils/shampoo_utils.py b/distributed_shampoo/utils/shampoo_utils.py
index 23060d9..8d071ad 100644
--- a/distributed_shampoo/utils/shampoo_utils.py
+++ b/distributed_shampoo/utils/shampoo_utils.py
@@ -7,6 +7,7 @@
 
 """
 
+import copy
 import heapq
 import logging
 import math
@@ -15,10 +16,16 @@
 from functools import cache, partial, reduce
 from itertools import accumulate, chain, compress, islice, pairwise
 from types import TracebackType
-from typing import Any, TypeVar
+from typing import Any, TypeGuard, TypeVar
 
 import torch
-from distributed_shampoo.shampoo_types import LoadBalancingConfig
+from distributed_shampoo.shampoo_types import (
+    DISTRIBUTED_CONFIG,
+    FullyShardDistributedConfig,
+    HybridShardDistributedConfig,
+    LoadBalancingConfig,
+    PARAMS,
+)
 from distributed_shampoo.utils.load_balancing_utils import DefaultCostModel
 from torch import distributed as dist, Tensor
 from torch.distributed.tensor import DTensor
@@ -420,3 +427,985 @@ def redistribute_and_update_params(
     torch._foreach_copy_(
         [p.to_local().flatten() for p in params], update_param_buffers[: len(params)]
     )
+
+
+def _compute_chunk_sizes(numel: int, num_chunks: int) -> list[int]:
+    """Compute chunk sizes that torch.chunk would produce.
+
+    torch.chunk(tensor, n, dim) semantics for n chunks along dim with size d:
+    - Computes chunk_size = ceil(d / n)
+    - First (n-1) chunks get chunk_size elements, last chunk gets remainder
+    - If d < n, only d chunks are produced (each of size 1)
+
+    Args:
+        numel (int): Size of the dimension to split.
+        num_chunks (int): Number of chunks to split into.
+
+    Returns:
+        sizes (list[int]): List of chunk sizes, one per chunk.
+    """
+    if numel == 0:
+        return [0] * num_chunks
+
+    if numel >= num_chunks:
+        # torch.chunk computes ceil(numel / num_chunks) as the chunk size
+        chunk_size = (numel + num_chunks - 1) // num_chunks  # This is ceil division
+        # First chunks get chunk_size, last chunk gets remainder
+        sizes = []
+        remaining = numel
+        for _i in range(num_chunks):
+            if remaining <= 0:
+                sizes.append(0)
+            elif remaining >= chunk_size:
+                sizes.append(chunk_size)
+                remaining -= chunk_size
+            else:
+                sizes.append(remaining)
+                remaining = 0
+        return sizes
+    else:
+        # numel < num_chunks: each element becomes its own chunk
+        # Ranks beyond numel get empty chunks (size 0)
+        sizes = [1] * numel + [0] * (num_chunks - numel)
+        return sizes
+
+
+def _compute_param_chunk_sizes(
+    params: tuple[DTensor, ...], group_size: int
+) -> list[list[int]]:
+    """Compute per-rank chunk sizes for all params based on FSDP dim-0 sharding.
+
+    For each parameter, computes how many elements each rank holds after
+    FSDP's dim-0 sharding (torch.chunk semantics).
+
+    Args:
+        params (tuple[DTensor, ...]): Tuple of DTensor parameters.
+        group_size (int): Number of ranks in the process group.
+
+    Returns:
+        param_chunk_sizes (list[list[int]]): List of lists, where param_chunk_sizes[i][r] is the
+            number of elements rank r holds for parameter i.
+    """
+    param_chunk_sizes: list[list[int]] = []
+    for param in params:
+        global_shape = tuple(param.shape)
+
+        if len(global_shape) == 0:
+            # Scalar tensor - all on rank 0
+            chunk_sizes = [1] + [0] * (group_size - 1)
+        else:
+            dim0_size = global_shape[0]
+            remaining_numel = math.prod(global_shape[1:])
+
+            # Compute dim-0 chunk sizes using torch.chunk semantics
+            dim0_chunks = _compute_chunk_sizes(dim0_size, group_size)
+            # Multiply by remaining dimensions to get actual element counts
+            chunk_sizes = [c * remaining_numel for c in dim0_chunks]
+
+        param_chunk_sizes.append(chunk_sizes)
+
+    return param_chunk_sizes
+
+
+def _build_buffer_views(buffer: Tensor, sizes: list[int]) -> list[Tensor]:
+    """Build contiguous views into a flat buffer given per-section sizes."""
+    views: list[Tensor] = []
+    offset = 0
+    for size in sizes:
+        views.append(buffer[offset : offset + size])
+        offset += size
+    return views
+
+
+class RedistributeParamsContext:
+    """Context for optimized parameter redistribution using all_to_all.
+
+    This class precomputes all static information needed for efficient parameter
+    redistribution during initialization, then uses that information to perform
+    fast redistribution using a single all_to_all collective instead of
+    multiple point-to-point calls.
+
+    The key insight is that torch.chunk(tensor, group_size, dim=0) is used to split
+    full params, and FSDP shards along dim=0. So chunk[i] corresponds to rank i's
+    local shard.
+
+    TODO(irisz): Currently assumes dim-0 sharding only. FSDP2 supports non-dim-0 sharding
+    but it is not enabled on APS. Extend to support arbitrary shard dimensions if needed.
+
+    Args:
+        params (tuple[DTensor, ...]): Tuple of DTensor parameters to redistribute.
+        assigned_params_mask (tuple[bool, ...]): Boolean mask indicating which params this rank computes.
+        dist_group (torch.distributed.ProcessGroup): Process group for communication.
+
+    Example:
+        # During initialization:
+        ctx = RedistributeParamsContext(params, assigned_mask, dist_group)
+
+        # During each optimizer step:
+        ctx.redistribute_and_update_params(local_full_params)
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        params: tuple[DTensor, ...],
+        assigned_params_mask: tuple[bool, ...],
+        dist_group: torch.distributed.ProcessGroup,
+    ) -> None:
+        self._params: tuple[DTensor, ...] = params
+        self._assigned_params_mask: tuple[bool, ...] = assigned_params_mask
+        self._dist_group: torch.distributed.ProcessGroup = dist_group
+        self._group_size: int = dist_group.size()
+        self._rank: int = dist.get_rank(group=dist_group)
+
+        if not params:
+            raise ValueError("params cannot be empty")
+
+        if len(params) != len(assigned_params_mask):
+            raise ValueError(
+                f"len(params) ({len(params)}) != "
+                f"len(assigned_params_mask) ({len(assigned_params_mask)})"
+            )
+
+        self._dtype: torch.dtype = params[0].to_local().dtype
+        self._device: torch.device = params[0].to_local().device
+
+        # Validate all params have same dtype
+        if any(p.to_local().dtype != self._dtype for p in params):
+            raise NotImplementedError(
+                "When using optimized redistribution, parameters of "
+                "different dtypes are not currently supported."
+            )
+
+        # Precompute static metadata
+        self._precompute_metadata()
+
+    def _precompute_metadata(self) -> None:
+        """Precompute all static information needed for redistribution.
+
+        Key insight: FSDP2 uses torch.chunk for dim-0 sharding. We can compute
+        chunk sizes mathematically during initialization.
+
+        torch.chunk(tensor, n) semantics:
+        - If numel >= n: returns n chunks, each of size ceil(numel/n) except last
+        - If numel < n: returns numel chunks of size 1 (fewer than n chunks!)
+        """
+        group_size = self._group_size
+
+        # Store the local shard numel for each param (what FSDP actually shards to)
+        self._param_local_numels: list[int] = [
+            p.to_local().numel() for p in self._params
+        ]
+        # Store the global numel for each param
+        self._param_global_numels: list[int] = [p.numel() for p in self._params]
+
+        # Compute which params each rank owns (based on round-robin assignment)
+        self._local_param_indices: list[int] = [
+            i for i, assigned in enumerate(self._assigned_params_mask) if assigned
+        ]
+
+        # Compute chunk sizes using FSDP's dim-0 sharding semantics
+        self._param_chunk_sizes: list[list[int]] = _compute_param_chunk_sizes(
+            self._params, group_size
+        )
+
+        # Compute send sizes: sum of chunk sizes for owned params going to each peer
+        self._send_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            send_size = sum(
+                self._param_chunk_sizes[param_idx][peer_rank]
+                for param_idx in self._local_param_indices
+            )
+            self._send_sizes.append(send_size)
+        self._total_send_size = sum(self._send_sizes)
+
+        # Compute recv sizes: sum of chunk sizes for params owned by each peer
+        self._recv_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            recv_size = sum(
+                self._param_chunk_sizes[param_idx][self._rank]
+                for param_idx in peer_param_indices
+            )
+            self._recv_sizes.append(recv_size)
+        self._total_recv_size = sum(self._recv_sizes)
+
+        # Pre-allocate persistent send and recv buffers for all_to_all.
+        # Memory overhead: O(total_params / world_size) each.
+        #   - send_buffer holds chunked assigned (owned) params = total_params / world_size
+        #   - recv_buffer holds received shards for all params from all peers,
+        #     but each param's local shard is ~1/world_size of its full size,
+        #     so total recv ≈ total_params / world_size as well.
+        self._send_buffer = torch.empty(
+            self._total_send_size, dtype=self._dtype, device=self._device
+        )
+        self._recv_buffer = torch.empty(
+            self._total_recv_size, dtype=self._dtype, device=self._device
+        )
+
+        # Pre-compute send_list and recv_list as contiguous views into buffers.
+        self._send_list = _build_buffer_views(self._send_buffer, self._send_sizes)
+        self._recv_list = _build_buffer_views(self._recv_buffer, self._recv_sizes)
+
+        # Pre-compute send and recv views for batched operations.
+        # Forward-declare attributes set by helper methods for Pyre.
+        self._send_dst_views: list[Tensor] = []
+        self._send_param_indices: list[int] = []
+        self._send_peer_ranks: list[int] = []
+        self._param_to_local_idx: dict[int, int] = {}
+        self._local_param_global_shapes: list[tuple[int, ...]] = []
+        self._param_recv_info: list[tuple[int, int]] = []
+        self._recv_src_views: list[Tensor] = []
+        self._recv_dst_views: list[Tensor] = []
+        self._compute_send_views()
+        self._compute_recv_views()
+
+        logger.info(
+            f"RedistributeParamsContext initialized: "
+            f"rank={self._rank}, group_size={group_size}, "
+            f"num_params={len(self._params)}, "
+            f"local_params={len(self._local_param_indices)}, "
+            f"send_sizes={self._send_sizes}, "
+            f"recv_sizes={self._recv_sizes}, "
+            f"param_local_numels={self._param_local_numels}, "
+            f"param_global_numels={self._param_global_numels}, "
+            f"param_chunk_sizes={self._param_chunk_sizes}"
+        )
+
+    def _compute_send_views(self) -> None:
+        """Pre-compute per-(local_param, peer_rank) destination views into send buffer."""
+        group_size = self._group_size
+        send_offset = 0
+        for peer_rank in range(group_size):
+            for param_idx in self._local_param_indices:
+                chunk_size = self._param_chunk_sizes[param_idx][peer_rank]
+                if chunk_size > 0:
+                    self._send_dst_views.append(
+                        self._send_buffer[send_offset : send_offset + chunk_size]
+                    )
+                    self._send_param_indices.append(param_idx)
+                    self._send_peer_ranks.append(peer_rank)
+                    send_offset += chunk_size
+
+        # Build a mapping from param_idx to local_idx for fast lookup.
+        self._param_to_local_idx.update(
+            {
+                param_idx: local_idx
+                for local_idx, param_idx in enumerate(self._local_param_indices)
+            }
+        )
+
+        # Pre-compute global shapes for owned params (used for chunking).
+        self._local_param_global_shapes.extend(
+            tuple(self._params[param_idx].shape)
+            for param_idx in self._local_param_indices
+        )
+
+    def _compute_recv_views(self) -> None:
+        """Pre-compute recv unpacking info and copy-out views."""
+        group_size = self._group_size
+
+        # Precompute recv unpacking info
+        # Recv buffer layout: [data from rank 0][data from rank 1]...
+        # _param_recv_info[param_idx] = (offset in recv_buffer, chunk_size)
+        self._param_recv_info = [(-1, -1)] * len(self._params)
+        recv_offset = 0
+        for peer_rank in range(group_size):
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            for param_idx in peer_param_indices:
+                chunk_size = self._param_chunk_sizes[param_idx][self._rank]
+                self._param_recv_info[param_idx] = (recv_offset, chunk_size)
+                recv_offset += chunk_size
+
+        # Pre-compute recv copy-out views for batched copy in
+        # redistribute_and_update_params().
+        for param_idx in range(len(self._params)):
+            recv_offset, chunk_size = self._param_recv_info[param_idx]
+            if chunk_size > 0:
+                local_param = self._params[param_idx].to_local()
+                local_numel = local_param.numel()
+                assert chunk_size == local_numel, (
+                    f"chunk_size ({chunk_size}) != local_numel ({local_numel}) "
+                    f"for param {param_idx}"
+                )
+                self._recv_src_views.append(
+                    self._recv_buffer[recv_offset : recv_offset + chunk_size]
+                )
+                self._recv_dst_views.append(local_param.flatten())
+
+    @torch.no_grad()
+    def redistribute_and_update_params(
+        self,
+        local_full_params: list[Tensor],
+    ) -> None:
+        """Redistribute updated parameters using a SINGLE coalesced all_to_all.
+
+        This combines ALL params into one collective call instead of multiple rounds,
+        reducing the number of collective operations.
+
+        Args:
+            local_full_params (list[Tensor]): List of full updated params computed by this rank.
+                Must match the params indicated by assigned_params_mask.
+        """
+        assert len(local_full_params) == len(self._local_param_indices), (
+            f"Expected {len(self._local_param_indices)} local params, "
+            f"got {len(local_full_params)}"
+        )
+
+        group_size = self._group_size
+
+        # Pre-chunk all local full params by peer rank for source view construction.
+        # param_chunks[local_idx][peer_rank] = flattened chunk tensor
+        param_chunks: list[list[Tensor]] = []
+        for local_idx in range(len(self._local_param_indices)):
+            full_param = local_full_params[local_idx]
+            global_shape = self._local_param_global_shapes[local_idx]
+
+            if len(global_shape) == 0:
+                # Scalar tensor - all goes to rank 0.
+                # Note: FSDP2 does not produce scalar params in practice;
+                # this branch is defensive.
+                chunks = [full_param.flatten()]
+                while len(chunks) < group_size:
+                    chunks.append(
+                        torch.empty(0, dtype=full_param.dtype, device=full_param.device)
+                    )
+            else:
+                full_param_reshaped = full_param.view(global_shape)
+                dim0_chunks = torch.chunk(full_param_reshaped, group_size, dim=0)
+                chunks = [c.flatten() for c in dim0_chunks]
+                while len(chunks) < group_size:
+                    chunks.append(
+                        torch.empty(0, dtype=full_param.dtype, device=full_param.device)
+                    )
+            param_chunks.append(chunks)
+
+        # Build source list matching the pre-computed destination views and
+        # batch-copy into the send buffer.
+        src_list: list[Tensor] = []
+        dst_list: list[Tensor] = []
+        for dst_view, param_idx, peer_rank in zip(
+            self._send_dst_views,
+            self._send_param_indices,
+            self._send_peer_ranks,
+            strict=True,
+        ):
+            local_idx = self._param_to_local_idx[param_idx]
+            src_list.append(param_chunks[local_idx][peer_rank])
+            dst_list.append(dst_view)
+
+        if dst_list:
+            torch._foreach_copy_(dst_list, src_list)
+
+        # Execute SINGLE all_to_all - this is the key optimization!
+        dist.all_to_all(self._recv_list, self._send_list, group=self._dist_group)
+
+        # Batch-copy received chunks to local param shards in one fused kernel.
+        if self._recv_dst_views:
+            torch._foreach_copy_(self._recv_dst_views, self._recv_src_views)
+
+
+class GatherGradientsContext:
+    """Context for gathering gradients to owning ranks using all_to_all.
+
+    This is the mirror of RedistributeParamsContext. While RedistributeParamsContext
+    sends full params from owning ranks to all ranks' local shards,
+    GatherGradientsContext gathers local gradient shards from all ranks to the
+    owning rank to reconstruct full gradients.
+
+    Data flow:
+        Each rank has local grad shards for ALL params (from FSDP).
+        Each rank sends its local shard of param i to the rank that owns param i.
+        The owning rank concatenates received shards along dim-0 to get the full gradient.
+
+    This replaces per-param full_tensor() all-gathers with a single all_to_all,
+    reducing peak memory from O(total_params) to O(total_params / world_size).
+
+    Args:
+        params (tuple[DTensor, ...]): Tuple of DTensor parameters.
+        assigned_params_mask (tuple[bool, ...]): Boolean mask indicating which params this rank owns.
+        dist_group (torch.distributed.ProcessGroup): Process group for communication.
+
+    Example:
+        # During initialization:
+        ctx = GatherGradientsContext(params, assigned_mask, dist_group)
+
+        # During each optimizer step:
+        full_grads = ctx.gather_gradients()
+    """
+
+    @torch.no_grad()
+    def __init__(
+        self,
+        params: tuple[DTensor, ...],
+        assigned_params_mask: tuple[bool, ...],
+        dist_group: torch.distributed.ProcessGroup,
+    ) -> None:
+        self._params = params
+        self._assigned_params_mask = assigned_params_mask
+        self._dist_group = dist_group
+        self._group_size: int = dist_group.size()
+        self._rank: int = dist.get_rank(group=dist_group)
+
+        if not params:
+            raise ValueError("params cannot be empty")
+
+        if len(params) != len(assigned_params_mask):
+            raise ValueError(
+                f"len(params) ({len(params)}) != "
+                f"len(assigned_params_mask) ({len(assigned_params_mask)})"
+            )
+
+        self._dtype: torch.dtype = params[0].to_local().dtype
+        self._device: torch.device = params[0].to_local().device
+
+        # Validate all params have same dtype
+        if any(p.to_local().dtype != self._dtype for p in params):
+            raise NotImplementedError(
+                "When using gradient gathering, parameters of "
+                "different dtypes are not currently supported."
+            )
+
+        self._precompute_metadata()
+
+    def _precompute_metadata(self) -> None:
+        """Precompute all static information needed for gradient gathering.
+
+        The send/recv directions are swapped relative to RedistributeParamsContext:
+        - Send: this rank sends its local grad shard of param i to the rank owning param i
+        - Recv: this rank receives grad shards from all ranks for params it owns
+
+        The send buffer layout (per peer rank) groups params by owner:
+          send_to_peer_r = [local_shard_of_param_j for j owned by peer r]
+
+        The recv buffer layout (per peer rank) groups by sender:
+          recv_from_peer_r = [shard_from_rank_r_for_param_j for j owned by this rank]
+        """
+        group_size = self._group_size
+
+        # Compute which params this rank owns
+        self._local_param_indices: list[int] = [
+            i for i, assigned in enumerate(self._assigned_params_mask) if assigned
+        ]
+
+        # Compute chunk sizes using shared utility
+        self._param_chunk_sizes: list[list[int]] = _compute_param_chunk_sizes(
+            self._params, group_size
+        )
+
+        # Store global shapes for assigned params (used to reconstruct full grads)
+        self._param_global_shapes: list[tuple[int, ...]] = [
+            tuple(p.shape) for p in self._params
+        ]
+
+        # Send sizes: for each peer rank, sum of this rank's local shard sizes
+        # for params owned by that peer.
+        # "This rank sends its chunk[self._rank] of param i to the rank that owns param i."
+        self._send_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            # Params owned by peer_rank
+            peer_param_indices = [
+                i for i in range(len(self._params)) if i % group_size == peer_rank
+            ]
+            send_size = sum(
+                self._param_chunk_sizes[param_idx][self._rank]
+                for param_idx in peer_param_indices
+            )
+            self._send_sizes.append(send_size)
+        self._total_send_size = sum(self._send_sizes)
+
+        # Recv sizes: for each peer rank, sum of that peer's shard sizes
+        # for params owned by this rank.
+        # "This rank receives chunk[peer_rank] of param j from peer_rank, for all j this rank owns."
+        self._recv_sizes: list[int] = []
+        for peer_rank in range(group_size):
+            recv_size = sum(
+                self._param_chunk_sizes[param_idx][peer_rank]
+                for param_idx in self._local_param_indices
+            )
+            self._recv_sizes.append(recv_size)
+        self._total_recv_size = sum(self._recv_sizes)
+
+        # Pre-allocate persistent send and recv buffers.
+        # Memory overhead: O(total_params / world_size) each.
+        #   - send_buffer holds local grad shards for all params ≈ total_params / world_size
+        #   - recv_buffer holds received grad shards for assigned (owned) params
+        #     from all peers = total_params / world_size
+        self._send_buffer = torch.empty(
+            self._total_send_size, dtype=self._dtype, device=self._device
+        )
+        self._recv_buffer = torch.empty(
+            self._total_recv_size, dtype=self._dtype, device=self._device
+        )
+
+        # Pre-compute send_list and recv_list as contiguous views into buffers.
+        self._send_list = _build_buffer_views(self._send_buffer, self._send_sizes)
+        self._recv_list = _build_buffer_views(self._recv_buffer, self._recv_sizes)
+
+        # Precompute recv unpacking info for each assigned param.
+        # After all_to_all, for assigned param j, its shards from each rank are
+        # scattered across the recv buffer (one piece per peer rank section).
+        # We precompute (offset, size) pairs for each (param, peer_rank) so we
+        # can efficiently gather them into a full gradient.
+        #
+        # _param_recv_info[local_idx] = list of (offset, size) for each peer rank
+        # where local_idx indexes into self._local_param_indices
+        self._param_recv_info: list[list[tuple[int, int]]] = []
+        # First compute offsets within each peer_rank's recv section
+        # recv_offsets_per_peer[peer_rank] tracks the running offset within
+        # the recv section for peer_rank
+        recv_section_starts: list[int] = []
+        offset = 0
+        for recv_size in self._recv_sizes:
+            recv_section_starts.append(offset)
+            offset += recv_size
+
+        # For each local param, compute where each peer's shard lands in recv_buffer
+        # The recv section for peer_rank contains shards in local_param_indices order
+        peer_offsets = list(recv_section_starts)  # mutable copy
+        # We need to iterate local_param_indices in order for each peer_rank
+        # to compute offsets correctly
+        self._param_local_idx_map: dict[int, int] = {
+            param_idx: local_idx
+            for local_idx, param_idx in enumerate(self._local_param_indices)
+        }
+        # Initialize recv info for each local param
+        for _ in self._local_param_indices:
+            self._param_recv_info.append([(-1, -1)] * group_size)
+
+        for peer_rank in range(group_size):
+            for param_idx in self._local_param_indices:
+                local_idx = self._param_local_idx_map[param_idx]
+                chunk_size = self._param_chunk_sizes[param_idx][peer_rank]
+                self._param_recv_info[local_idx][peer_rank] = (
+                    peer_offsets[peer_rank],
+                    chunk_size,
+                )
+                peer_offsets[peer_rank] += chunk_size
+
+        # Pre-compute views for batched operations.
+        # Forward-declare attributes set by helper methods for Pyre.
+        self._full_grad_buffers: list[Tensor] = []
+        self._unpack_src_views: list[Tensor] = []
+        self._unpack_dst_views: list[Tensor] = []
+        self._send_dst_views: list[Tensor] = []
+        self._send_src_param_indices: list[int] = []
+        self._send_src_chunk_sizes: list[int] = []
+        self._compute_unpack_views()
+        self._compute_send_dst_views()
+
+        logger.info(
+            f"GatherGradientsContext initialized: "
+            f"rank={self._rank}, group_size={group_size}, "
+            f"num_params={len(self._params)}, "
+            f"local_params={len(self._local_param_indices)}, "
+            f"send_sizes={self._send_sizes}, "
+            f"recv_sizes={self._recv_sizes}"
+        )
+
+    def _compute_unpack_views(self) -> None:
+        """Pre-allocate full gradient buffers and pre-compute recv unpack views."""
+        for local_idx in range(len(self._local_param_indices)):
+            total_size = sum(
+                size for _, size in self._param_recv_info[local_idx] if size > 0
+            )
+            grad_buffer = torch.empty(
+                total_size, dtype=self._dtype, device=self._device
+            )
+            self._full_grad_buffers.append(grad_buffer)
+            dst_offset = 0
+            for recv_offset, size in self._param_recv_info[local_idx]:
+                if size > 0:
+                    self._unpack_src_views.append(
+                        self._recv_buffer[recv_offset : recv_offset + size]
+                    )
+                    self._unpack_dst_views.append(
+                        grad_buffer[dst_offset : dst_offset + size]
+                    )
+                    dst_offset += size
+
+    def _compute_send_dst_views(self) -> None:
+        """Pre-compute destination views into the send buffer for batched copy."""
+        group_size = self._group_size
+        send_offset = 0
+        for peer_rank in range(group_size):
+            for param_idx in range(len(self._params)):
+                if param_idx % group_size != peer_rank:
+                    continue
+                chunk_size = self._param_chunk_sizes[param_idx][self._rank]
+                if chunk_size > 0:
+                    self._send_dst_views.append(
+                        self._send_buffer[send_offset : send_offset + chunk_size]
+                    )
+                    self._send_src_param_indices.append(param_idx)
+                    self._send_src_chunk_sizes.append(chunk_size)
+                    send_offset += chunk_size
+
+    def _build_send_list(self, has_grad: list[bool]) -> None:
+        """Pack local grad shards into the persistent send buffer.
+
+        For each pre-computed (param_idx, dst_view) pair, copies the local
+        gradient shard into the send buffer using foreach_copy_ for efficiency.
+
+        Args:
+            has_grad: Boolean mask indicating which params have gradients.
+        """
+        src_list: list[Tensor] = []
+        dst_list: list[Tensor] = []
+
+        for dst, param_idx, chunk_size in zip(
+            self._send_dst_views,
+            self._send_src_param_indices,
+            self._send_src_chunk_sizes,
+            strict=True,
+        ):
+            if has_grad[param_idx]:
+                local_grad = self._params[param_idx].grad
+                assert local_grad is not None
+                src_list.append(
+                    local_grad.to_local().flatten()[:chunk_size]  # type: ignore
+                )
+                dst_list.append(dst)
+            else:
+                dst.zero_()
+
+        if dst_list:
+            torch._foreach_copy_(dst_list, src_list)
+
+    def _unpack_recv_buffer(self, has_grad: list[bool]) -> list[Tensor | None]:
+        """Unpack the recv buffer into full gradients for assigned params.
+
+        Uses pre-computed views and foreach_copy_ to batch-copy all shards from
+        the recv buffer into pre-allocated full gradient buffers, then reshapes
+        each to the original global shape.
+
+        Args:
+            has_grad (list[bool]): Boolean mask indicating which params have gradients.
+
+        Returns:
+            full_grads (list[Tensor | None]): List of full gradients (same length as self._params).
+                None for params without gradients or unassigned params.
+        """
+        # Batch-copy all recv shards into pre-allocated full grad buffers.
+        if self._unpack_dst_views:
+            torch._foreach_copy_(self._unpack_dst_views, self._unpack_src_views)
+
+        full_grads: list[Tensor | None] = [None for _ in self._params]
+        for local_idx, param_idx in enumerate(self._local_param_indices):
+            if not has_grad[param_idx]:
+                continue
+            global_shape = self._param_global_shapes[param_idx]
+            if len(global_shape) == 0:
+                full_grads[param_idx] = self._full_grad_buffers[local_idx].squeeze()
+            else:
+                full_grads[param_idx] = self._full_grad_buffers[local_idx].view(
+                    global_shape
+                )
+
+        return full_grads
+
+    def _execute_all_to_all(self) -> None:
+        """Execute all_to_all using persistent send/recv buffers."""
+        dist.all_to_all(self._recv_list, self._send_list, group=self._dist_group)
+
+    @torch.no_grad()
+    def gather_gradients(self) -> list[Tensor | None]:
+        """Gather gradients from all ranks using a single all_to_all.
+
+        Each rank sends its local grad shard of each param to the rank that owns
+        that param. The owning rank then concatenates the received shards along
+        dim-0 to reconstruct the full gradient.
+
+        Returns:
+            full_grads: List of full gradients for ALL params (same length as
+                self._params). Entries are None for params not assigned to this
+                rank, regardless of whether they have gradients. Only params
+                assigned to this rank will have non-None values.
+        """
+        has_grad = [p.grad is not None for p in self._params]
+
+        self._build_send_list(has_grad)
+        self._execute_all_to_all()
+
+        result = self._unpack_recv_buffer(has_grad)
+        # Clone to detach from _full_grad_buffers, which will be overwritten
+        # by subsequent gather_gradients() calls.
+        return [t.clone() if t is not None else None for t in result]
+
+    @torch.no_grad()
+    def gather_params(self) -> list[Tensor | None]:
+        """Gather full parameter values from all ranks using a single all_to_all.
+
+        Similar to gather_gradients(), but operates on parameter values instead
+        of gradients. Every parameter always has a value, so no None handling
+        is needed on the send side.
+
+        Returns:
+            full_params: List of full parameter tensors (same length as
+                self._params). Entries are None for params not assigned to this
+                rank. Only params assigned to this rank will have non-None values.
+        """
+        # Pack param values into persistent send buffer using foreach_copy_.
+        src_list: list[Tensor] = []
+        dst_list: list[Tensor] = []
+
+        for dst, param_idx, chunk_size in zip(
+            self._send_dst_views,
+            self._send_src_param_indices,
+            self._send_src_chunk_sizes,
+            strict=True,
+        ):
+            src_list.append(self._params[param_idx].to_local().flatten()[:chunk_size])
+            dst_list.append(dst)
+
+        if dst_list:
+            torch._foreach_copy_(dst_list, src_list)
+
+        self._execute_all_to_all()
+
+        result = self._unpack_recv_buffer(has_grad=[True] * len(self._params))
+        # Clone to detach from _full_grad_buffers, which will be overwritten
+        # by subsequent gather_gradients() calls.
+        return [t.clone() if t is not None else None for t in result]
+
+
+def _device_key(device: torch.device) -> tuple[str, int]:
+    """Canonicalizes a torch.device into a hashable (type, index) tuple.
+
+    Resolves cuda devices without an explicit index to the current device, so
+    torch.device('cuda') and torch.device('cuda:0') normalize to the same key
+    and don't produce duplicate cache entries in `_get_triu_indices`.
+    """
+    if device.type == "cuda" and device.index is None:
+        return ("cuda", torch.cuda.current_device())
+    return (device.type, device.index or 0)
+
+
+@cache
+def _get_triu_indices(d: int, device_type: str, device_index: int) -> Tensor:
+    """Returns cached upper-triangular index pairs as a (2, d*(d+1)/2) int32 tensor.
+
+    Args:
+        d (int): Dimension of the square matrix.
+        device_type (str): Device type, e.g. "cpu" or "cuda".
+        device_index (int): Resolved device index. Use `_device_key` to
+            normalize from a torch.device.
+
+    Returns:
+        indices (Tensor): A (2, d*(d+1)/2) int32 tensor of (row, col) index
+            pairs for the upper triangle in row-major order. The returned
+            tensor is cached by reference and MUST NOT be mutated by callers.
+
+    Notes:
+        Cached process-wide via functools.cache. int32 is sufficient since
+        matrix dim d is bounded by max_preconditioner_dim (default 1024, up
+        to 4096 in known production workloads), well under int32's 2**31
+        range; using int32 instead of torch.triu_indices' int64 default
+        halves the cache footprint.
+    """
+    return torch.triu_indices(d, d, device=torch.device(device_type, device_index)).to(
+        torch.int32
+    )
+
+
+def pack_upper_triangular(matrix: Tensor) -> Tensor:
+    """Packs the upper triangle of a symmetric matrix into a flat 1D tensor.
+
+    For a (d, d) symmetric matrix, extracts the upper triangular elements
+    (including diagonal) into a contiguous 1D tensor of size d*(d+1)/2.
+
+    Args:
+        matrix (Tensor): A 2D symmetric matrix of shape (d, d).
+
+    Returns:
+        packed (Tensor): A 1D tensor of size d*(d+1)/2 containing the upper
+            triangular elements in row-major order.
+    """
+    rows, cols = _get_triu_indices(matrix.shape[0], *_device_key(matrix.device))
+    return matrix[rows, cols].contiguous()
+
+
+def unpack_upper_triangular(packed: Tensor, dim: int) -> Tensor:
+    """Reconstructs a symmetric matrix from its packed upper triangle.
+
+    Args:
+        packed (Tensor): A 1D tensor of size dim*(dim+1)/2 containing the
+            upper triangular elements in row-major order.
+        dim (int): The dimension of the original square matrix.
+
+    Returns:
+        matrix (Tensor): A 2D symmetric matrix of shape (dim, dim).
+    """
+    matrix = torch.zeros(dim, dim, dtype=packed.dtype, device=packed.device)
+    rows, cols = _get_triu_indices(dim, *_device_key(packed.device))
+    matrix[rows, cols] = packed  # Fill upper triangle.
+    matrix[cols, rows] = packed  # Mirror to lower triangle (swap row/col indices).
+    return matrix
+
+
+_T = TypeVar("_T")
+
+
+def greedy_bin_pack(
+    items: Sequence[_T],
+    num_bins: int,
+    cost_fn: Callable[[_T], int],
+) -> tuple[list[list[_T]], list[int]]:
+    """Partition items into bins using greedy LPT (Longest Processing Time) bin-packing.
+
+    Sorts items by cost (largest first) and assigns each to the bin with the
+    smallest total cost. This balances total cost across bins.
+
+    Args:
+        items (Sequence[_T]): Items to partition.
+        num_bins (int): Number of bins.
+        cost_fn (Callable[[_T], int]): Function returning the cost of an item.
+
+    Returns:
+        A tuple of (bins, bin_costs) where bins[i] is the list of items in bin i
+        and bin_costs[i] is the total cost of bin i. Empty bins are included.
+    """
+    bins: list[list[_T]] = [[] for _ in range(num_bins)]
+    heap: list[tuple[int, int]] = [(0, i) for i in range(num_bins)]
+    heapq.heapify(heap)
+    # Use enumerate as tiebreaker to avoid comparing items that may not
+    # support < (e.g., torch.nn.Parameter).
+    for cost, _, item in sorted(
+        ((cost_fn(item), i, item) for i, item in enumerate(items)),
+        reverse=True,
+    ):
+        min_cost, min_idx = heapq.heappop(heap)
+        bins[min_idx].append(item)
+        heapq.heappush(heap, (min_cost + cost, min_idx))
+    bin_costs = [0] * num_bins
+    for cost, idx in heap:
+        bin_costs[idx] = cost
+    return bins, bin_costs
+
+
+def split_param_groups(
+    param_groups: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    """Split each FSDP/HSDP lossless param group with ``num_sub_groups > 1``
+    into N independent sub-groups, each with its own copy of the distributed
+    config (and thus its own distributor + NCCL communicator).
+
+    Sub-groups can then run on dedicated CUDA streams in parallel threads
+    (see ``DistributedShampoo._run_threaded_step``) to overlap compute and
+    communication.
+
+    Within a qualifying group, params are distributed across sub-groups via
+    greedy bin-packing on numel to balance element count.
+
+    Currently scoped to ``FullyShardDistributedConfig`` (and its subclass
+    ``HybridShardDistributedConfig``); other distributor types pass through
+    unchanged.
+
+    Args:
+        param_groups (list[dict[str, Any]]): Optimizer param groups to split.
+
+    Returns:
+        new_param_groups (list[dict[str, Any]]): New param groups list. Groups
+            that don't opt into splitting pass through; qualifying groups are
+            replaced by N sub-groups with the params evenly distributed.
+
+    TODO(irisz): Support automatic determination of optimal num_sub_groups
+    based on param count, compute cost, and communicator memory overhead.
+    """
+    new_param_groups: list[dict[str, Any]] = []
+    for group in param_groups:
+        distributed_config = group[DISTRIBUTED_CONFIG]
+        if not _should_split(distributed_config):
+            new_param_groups.append(group)
+            continue
+        _validate_num_sub_groups(distributed_config, len(group[PARAMS]))
+        sub_groups = _create_sub_groups(group, distributed_config.num_sub_groups)
+        new_param_groups.extend(sub_groups)
+        logger.info(
+            "Split param group (%d params) into %d sub-groups (numel per group: %s).",
+            len(group[PARAMS]),
+            len(sub_groups),
+            [sum(p.numel() for p in g[PARAMS]) for g in sub_groups],
+        )
+    return new_param_groups
+
+
+def _should_split(
+    distributed_config: Any,
+) -> TypeGuard[FullyShardDistributedConfig]:
+    """True when the config opts into sub-group splitting (num_sub_groups > 1)."""
+    return (
+        isinstance(distributed_config, FullyShardDistributedConfig)
+        and distributed_config.num_sub_groups > 1
+    )
+
+
+def _validate_num_sub_groups(
+    distributed_config: FullyShardDistributedConfig, num_params: int
+) -> None:
+    """Raise ValueError if num_sub_groups can't be satisfied for this group's params.
+
+    Each sub-group needs at least ``shard_size`` params so ROUND_ROBIN assigns
+    at least one param to each shard. Otherwise, shards with zero params deadlock
+    on all-to-all. For FSDP this collapses to ``shard_size=1`` (i.e., n <= num_params).
+    """
+    n = distributed_config.num_sub_groups
+    shard_size = (
+        distributed_config.device_mesh.size(1)
+        if isinstance(distributed_config, HybridShardDistributedConfig)
+        else 1
+    )
+    max_groups = max(1, num_params // shard_size)
+    if n > max_groups:
+        raise ValueError(
+            f"num_sub_groups={n} is too large for {num_params} params with "
+            f"shard_group_size={shard_size}. Each sub-group needs at least "
+            f"{shard_size} params for ROUND_ROBIN assignment. "
+            f"Maximum num_sub_groups={max_groups}."
+        )
+
+
+def _create_sub_groups(group: dict[str, Any], n: int) -> list[dict[str, Any]]:
+    """Split ``group``'s params into N sub-groups via greedy bin-packing on
+    numel; each carries a fresh copy of the distributed config so each gets
+    its own distributor and NCCL communicator.
+
+    Empty bins are dropped (can occur when one param's numel dominates the
+    total enough to leave a bin unfilled).
+
+    For HSDP, additionally validates that every non-empty bin holds at least
+    ``shard_size`` params. Greedy LPT can place a single dominant-numel param
+    alone in a bin, leaving fewer than ``shard_size`` params there even though
+    ``_validate_num_sub_groups`` (which assumes contiguous chunking) accepted
+    the count. ROUND_ROBIN within such a bin would assign zero params to some
+    shard ranks, deadlocking on all-to-all.
+    """
+    params = group[PARAMS]
+    distributed_config = group[DISTRIBUTED_CONFIG]
+    base_keys = {k: v for k, v in group.items() if k != PARAMS}
+    buckets, _ = greedy_bin_pack(items=params, num_bins=n, cost_fn=lambda p: p.numel())
+    if isinstance(distributed_config, HybridShardDistributedConfig):
+        shard_size = distributed_config.device_mesh.size(1)
+        for i, bucket in enumerate(buckets):
+            if bucket and len(bucket) < shard_size:
+                raise ValueError(
+                    f"HSDP sub-group {i} got {len(bucket)} params after "
+                    f"greedy bin-packing, fewer than shard_size={shard_size}. "
+                    f"This happens when one or more params dominate by numel; "
+                    f"ROUND_ROBIN would leave shard ranks with no params and "
+                    f"deadlock on all-to-all. Reduce num_sub_groups or "
+                    f"rebalance the param shapes."
+                )
+    return [
+        {
+            **base_keys,
+            PARAMS: bucket,
+            DISTRIBUTED_CONFIG: copy.copy(distributed_config),
+        }
+        for bucket in buckets
+        if bucket
+    ]
diff --git a/distributed_shampoo/utils/tests/abstract_dataclass_test.py b/distributed_shampoo/utils/tests/abstract_dataclass_test.py
index 46e9829..5a71a6d 100644
--- a/distributed_shampoo/utils/tests/abstract_dataclass_test.py
+++ b/distributed_shampoo/utils/tests/abstract_dataclass_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 from dataclasses import dataclass
diff --git a/distributed_shampoo/utils/tests/commons_test.py b/distributed_shampoo/utils/tests/commons_test.py
index 4b2eb1f..e67a5f6 100644
--- a/distributed_shampoo/utils/tests/commons_test.py
+++ b/distributed_shampoo/utils/tests/commons_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 from abc import ABC, abstractmethod
@@ -25,25 +23,25 @@ class BatchedTest(unittest.TestCase):
     def test_normal_batching(self) -> None:
         """Test batching an iterable with size divisible by batch size."""
         data = [1, 2, 3, 4, 5, 6]
-        result = list(batched(data, n=2))
+        result = list(batched(data, n=2))  # noqa: B911
         self.assertEqual(result, [(1, 2), (3, 4), (5, 6)])
 
     def test_uneven_batching(self) -> None:
         """Test batching an iterable with size not divisible by batch size."""
         data = [1, 2, 3, 4, 5]
-        result = list(batched(data, n=2))
+        result = list(batched(data, n=2))  # noqa: B911
         self.assertEqual(result, [(1, 2), (3, 4), (5,)])
 
     def test_empty_iterable(self) -> None:
         """Test batching an empty iterable."""
         data: list[int] = []
-        result = list(batched(data, n=3))
+        result = list(batched(data, n=3))  # noqa: B911
         self.assertEqual(result, [])
 
     def test_batch_size_one(self) -> None:
         """Test batching with batch size of 1."""
         data = [1, 2, 3]
-        result = list(batched(data, n=1))
+        result = list(batched(data, n=1))  # noqa: B911
         self.assertEqual(result, [(1,), (2,), (3,)])
 
     @parametrize("n", (-1, 0))
@@ -53,7 +51,7 @@ def test_invalid_batch_size(self, n: int) -> None:
         with self.assertRaisesRegex(
             ValueError, re.escape(f"{n=} must be at least one")
         ):
-            list(batched(data, n=n))
+            list(batched(data, n=n))  # noqa: B911
 
 
 class DummyRootClass:
diff --git a/distributed_shampoo/utils/tests/dict_zip_iterator_test.py b/distributed_shampoo/utils/tests/dict_zip_iterator_test.py
index 0a95d30..6522a82 100644
--- a/distributed_shampoo/utils/tests/dict_zip_iterator_test.py
+++ b/distributed_shampoo/utils/tests/dict_zip_iterator_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from collections.abc import Iterator
 
diff --git a/distributed_shampoo/utils/tests/optimizer_modules_test.py b/distributed_shampoo/utils/tests/optimizer_modules_test.py
index 6774c66..7255390 100644
--- a/distributed_shampoo/utils/tests/optimizer_modules_test.py
+++ b/distributed_shampoo/utils/tests/optimizer_modules_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from dataclasses import dataclass
 
diff --git a/distributed_shampoo/utils/tests/shampoo_model_utils_test.py b/distributed_shampoo/utils/tests/shampoo_model_utils_test.py
index 7c4edcc..3379cb8 100644
--- a/distributed_shampoo/utils/tests/shampoo_model_utils_test.py
+++ b/distributed_shampoo/utils/tests/shampoo_model_utils_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import unittest
 from math import sqrt
 from typing import cast
diff --git a/distributed_shampoo/utils/tests/shampoo_quantization_test.py b/distributed_shampoo/utils/tests/shampoo_quantization_test.py
index b128300..430d820 100644
--- a/distributed_shampoo/utils/tests/shampoo_quantization_test.py
+++ b/distributed_shampoo/utils/tests/shampoo_quantization_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import copy
 import re
 import unittest
diff --git a/distributed_shampoo/utils/tests/shampoo_state_dict_utils_test.py b/distributed_shampoo/utils/tests/shampoo_state_dict_utils_test.py
index 71c4d31..37c2816 100644
--- a/distributed_shampoo/utils/tests/shampoo_state_dict_utils_test.py
+++ b/distributed_shampoo/utils/tests/shampoo_state_dict_utils_test.py
@@ -7,8 +7,6 @@
 
 """
 
-#!/usr/bin/env python3
-
 import re
 import unittest
 
diff --git a/distributed_shampoo/utils/tests/shampoo_utils_test.py b/distributed_shampoo/utils/tests/shampoo_utils_test.py
index 484a8af..09fb40e 100644
--- a/distributed_shampoo/utils/tests/shampoo_utils_test.py
+++ b/distributed_shampoo/utils/tests/shampoo_utils_test.py
@@ -7,32 +7,50 @@
 
 """
 
-#!/usr/bin/env python3
-
 import math
 import re
 import unittest
 from operator import methodcaller
+from unittest.mock import MagicMock
 
 import torch
-from distributed_shampoo.shampoo_types import LoadBalancingConfig
+from distributed_shampoo.shampoo_types import (
+    DISTRIBUTED_CONFIG,
+    FSDPParamAssignmentStrategy,
+    FullyShardDistributedConfig,
+    HybridShardDistributedConfig,
+    LoadBalancingConfig,
+    PARAMS,
+    SingleDeviceDistributedConfig,
+)
 from distributed_shampoo.utils.load_balancing_utils import (
     PolynomialComputationalCostModel,
 )
+from distributed_shampoo.utils.shampoo_fully_shard_utils import _compute_chunk_sizes
 from distributed_shampoo.utils.shampoo_utils import (
+    _device_key,
+    _get_triu_indices,
     compress_list,
     distribute_buffer_sizes,
     generate_pairwise_indices,
     get_dtype_size,
+    greedy_bin_pack,
     merge_small_dims,
     multi_dim_split,
+    pack_upper_triangular,
     ParameterizeEnterExitContext,
+    split_param_groups,
+    unpack_upper_triangular,
 )
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
 )
 
+# Hoisted as an annotated module constant so pyre can infer the type of the
+# argument passed to @unittest.skipUnless (it cannot infer from a bare call).
+_HAS_CUDA: bool = torch.cuda.is_available()
+
 
 @instantiate_parametrized_tests
 class MergeSmallDimsTest(unittest.TestCase):
@@ -238,6 +256,69 @@ def test_generate_pairwise_indices_with_empty_list(self) -> None:
         )
 
 
+@instantiate_parametrized_tests
+class ComputeChunkSizesTest(unittest.TestCase):
+    """Tests for _compute_chunk_sizes which mirrors torch.chunk semantics."""
+
+    @parametrize(
+        "numel, num_chunks, expected",
+        [
+            # numel == 0: all chunks get size 0
+            (0, 1, [0]),
+            (0, 4, [0, 0, 0, 0]),
+            # numel >= num_chunks, evenly divisible
+            (4, 4, [1, 1, 1, 1]),
+            (8, 4, [2, 2, 2, 2]),
+            (12, 4, [3, 3, 3, 3]),
+            # numel >= num_chunks, not evenly divisible (last chunk gets remainder)
+            (10, 4, [3, 3, 3, 1]),
+            (7, 3, [3, 3, 1]),
+            (5, 2, [3, 2]),
+            (1, 1, [1]),
+            # numel < num_chunks: each element is its own chunk, extras are 0
+            (2, 4, [1, 1, 0, 0]),
+            (1, 4, [1, 0, 0, 0]),
+            (3, 5, [1, 1, 1, 0, 0]),
+        ],
+    )
+    def test_compute_chunk_sizes(
+        self, numel: int, num_chunks: int, expected: list[int]
+    ) -> None:
+        result = _compute_chunk_sizes(numel, num_chunks)
+        self.assertEqual(result, expected)
+
+    @parametrize("numel", (1, 5, 10, 13, 32, 100))
+    @parametrize("num_chunks", (1, 2, 3, 4, 7, 8))
+    def test_matches_torch_chunk(self, numel: int, num_chunks: int) -> None:
+        """Verify _compute_chunk_sizes matches actual torch.chunk behavior."""
+        tensor = torch.arange(numel)
+        actual_chunks = torch.chunk(tensor, num_chunks, dim=0)
+        actual_sizes = [c.numel() for c in actual_chunks]
+
+        computed_sizes = _compute_chunk_sizes(numel, num_chunks)
+        # torch.chunk may produce fewer chunks than num_chunks if numel < num_chunks.
+        # _compute_chunk_sizes pads with zeros in that case.
+        nonzero_computed = [s for s in computed_sizes if s > 0]
+        self.assertEqual(nonzero_computed, actual_sizes)
+        self.assertEqual(sum(computed_sizes), numel)
+
+    def test_sum_equals_numel(self) -> None:
+        """Verify all chunk sizes sum to the original numel."""
+        for numel in range(0, 20):
+            for num_chunks in range(1, 8):
+                sizes = _compute_chunk_sizes(numel, num_chunks)
+                self.assertEqual(
+                    sum(sizes),
+                    numel,
+                    f"Sum mismatch for numel={numel}, num_chunks={num_chunks}",
+                )
+                self.assertEqual(
+                    len(sizes),
+                    num_chunks,
+                    f"Length mismatch for numel={numel}, num_chunks={num_chunks}",
+                )
+
+
 class ParameterizeEnterExitContextTest(unittest.TestCase):
     """Test suite for the ParameterizeEnterExitContext class.
 
@@ -374,3 +455,297 @@ def test_distribute_buffer_sizes(
             ),
             expected_result,
         )
+
+
+@instantiate_parametrized_tests
+class PackUpperTriangularTest(unittest.TestCase):
+    def test_pack_2x2(self) -> None:
+        # [[1, 2], [3, 4]] in row-major upper-triangular order: (0,0), (0,1), (1,1).
+        torch.testing.assert_close(
+            pack_upper_triangular(torch.tensor([[1.0, 2.0], [3.0, 4.0]])),
+            torch.tensor([1.0, 2.0, 4.0]),
+        )
+
+    def test_pack_3x3(self) -> None:
+        # Row-major upper-triangular order for 3x3:
+        # (0,0), (0,1), (0,2), (1,1), (1,2), (2,2).
+        torch.testing.assert_close(
+            pack_upper_triangular(
+                torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+            ),
+            torch.tensor([1.0, 2.0, 3.0, 5.0, 6.0, 9.0]),
+        )
+
+    def test_pack_1x1(self) -> None:
+        torch.testing.assert_close(
+            pack_upper_triangular(torch.tensor([[42.0]])),
+            torch.tensor([42.0]),
+        )
+
+    @parametrize("dim", (1, 2, 3, 5, 8))
+    def test_pack_size_invariant(self, dim: int) -> None:
+        packed = pack_upper_triangular(torch.zeros(dim, dim))
+        self.assertEqual(packed.shape, (dim * (dim + 1) // 2,))
+
+    def test_pack_preserves_dtype_and_device(self) -> None:
+        for dtype in (torch.float32, torch.float64, torch.bfloat16):
+            packed = pack_upper_triangular(torch.zeros(3, 3, dtype=dtype))
+            self.assertEqual(packed.dtype, dtype)
+
+
+@instantiate_parametrized_tests
+class UnpackUpperTriangularTest(unittest.TestCase):
+    def test_unpack_2x2(self) -> None:
+        # [1, 2, 4] reconstructs to symmetric matrix [[1, 2], [2, 4]].
+        torch.testing.assert_close(
+            unpack_upper_triangular(torch.tensor([1.0, 2.0, 4.0]), dim=2),
+            torch.tensor([[1.0, 2.0], [2.0, 4.0]]),
+        )
+
+    def test_unpack_3x3(self) -> None:
+        torch.testing.assert_close(
+            unpack_upper_triangular(
+                torch.tensor([1.0, 2.0, 3.0, 5.0, 6.0, 9.0]), dim=3
+            ),
+            torch.tensor([[1.0, 2.0, 3.0], [2.0, 5.0, 6.0], [3.0, 6.0, 9.0]]),
+        )
+
+    def test_unpack_returns_symmetric(self) -> None:
+        unpacked = unpack_upper_triangular(
+            torch.arange(1, 11, dtype=torch.float32), dim=4
+        )
+        torch.testing.assert_close(unpacked, unpacked.T)
+
+    @parametrize("dim", (1, 2, 3, 5, 8))
+    def test_pack_unpack_roundtrip_symmetric(self, dim: int) -> None:
+        # Build a random symmetric matrix and verify pack -> unpack returns it.
+        a = torch.randn(dim, dim)
+        symmetric = a + a.T
+        torch.testing.assert_close(
+            unpack_upper_triangular(pack_upper_triangular(symmetric), dim=dim),
+            symmetric,
+        )
+
+    def test_unpack_preserves_dtype_and_device(self) -> None:
+        for dtype in (torch.float32, torch.float64, torch.bfloat16):
+            unpacked = unpack_upper_triangular(torch.zeros(6, dtype=dtype), dim=3)
+            self.assertEqual(unpacked.dtype, dtype)
+
+
+class TriuIndicesCacheTest(unittest.TestCase):
+    # The @cache on _get_triu_indices is process-wide; isolate every test
+    # so assertions on cache_info are independent of test ordering.
+    def setUp(self) -> None:
+        _get_triu_indices.cache_clear()
+
+    def tearDown(self) -> None:
+        _get_triu_indices.cache_clear()
+
+    def test_cache_reuse(self) -> None:
+        # First call populates the cache; subsequent same-dim calls hit it.
+        symmetric = torch.randn(8, 8)
+        symmetric = symmetric + symmetric.T
+        pack_upper_triangular(symmetric)
+        info_after_first = _get_triu_indices.cache_info()
+        self.assertEqual(info_after_first.misses, 1)
+        self.assertEqual(info_after_first.hits, 0)
+
+        for _ in range(5):
+            pack_upper_triangular(symmetric)
+            unpack_upper_triangular(torch.zeros(8 * 9 // 2), dim=8)
+        info_after_repeat = _get_triu_indices.cache_info()
+        self.assertEqual(info_after_repeat.misses, 1)
+        # 5 iterations * (1 pack + 1 unpack) = 10 hits, exact.
+        self.assertEqual(info_after_repeat.hits, 10)
+
+        # A different dim should miss exactly once more.
+        larger = torch.randn(16, 16)
+        larger = larger + larger.T
+        pack_upper_triangular(larger)
+        self.assertEqual(_get_triu_indices.cache_info().misses, 2)
+
+    def test_dtype_and_shape_and_values(self) -> None:
+        # int32 halves the cache footprint vs torch.triu_indices' int64 default.
+        # Lock dtype, shape, AND values against the reference torch.triu_indices.
+        indices = _get_triu_indices(32, "cpu", 0)
+        self.assertEqual(indices.dtype, torch.int32)
+        self.assertEqual(indices.shape, (2, 32 * 33 // 2))
+        torch.testing.assert_close(
+            indices, torch.triu_indices(32, 32, device="cpu").to(torch.int32)
+        )
+
+    def test_device_key_normalizes_cpu(self) -> None:
+        # cpu always normalizes to ('cpu', 0).
+        self.assertEqual(_device_key(torch.device("cpu")), ("cpu", 0))
+
+    @unittest.skipUnless(_HAS_CUDA, "CUDA required")
+    def test_device_key_normalizes_cuda(self) -> None:
+        # cuda without explicit index must collapse to the current device,
+        # so str-form aliases ('cuda' vs 'cuda:0') don't double-cache.
+        self.assertEqual(
+            _device_key(torch.device("cuda")),
+            _device_key(torch.device("cuda", torch.cuda.current_device())),
+        )
+
+    @unittest.skipUnless(_HAS_CUDA, "CUDA required")
+    def test_cuda_roundtrip(self) -> None:
+        symmetric = torch.randn(16, 16, device="cuda")
+        symmetric = symmetric + symmetric.T
+        packed = pack_upper_triangular(symmetric)
+        self.assertEqual(packed.device.type, "cuda")
+        unpacked = unpack_upper_triangular(packed, dim=16)
+        torch.testing.assert_close(unpacked, symmetric)
+
+
+@instantiate_parametrized_tests
+class SplitParamGroupsTest(unittest.TestCase):
+    @staticmethod
+    def _make_group(
+        num_params: int,
+        distributed_config: object,
+        extra: dict[str, object] | None = None,
+    ) -> dict[str, object]:
+        params = [torch.zeros(1) for _ in range(num_params)]
+        group: dict[str, object] = {
+            PARAMS: params,
+            DISTRIBUTED_CONFIG: distributed_config,
+        }
+        if extra is not None:
+            group.update(extra)
+        return group
+
+    @staticmethod
+    def _make_fsdp_config(num_sub_groups: int) -> FullyShardDistributedConfig:
+        return FullyShardDistributedConfig(
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=num_sub_groups,
+        )
+
+    def test_passthrough_when_num_sub_groups_is_1(self) -> None:
+        config = self._make_fsdp_config(num_sub_groups=1)
+        group = self._make_group(num_params=4, distributed_config=config)
+        result = split_param_groups([group])
+        self.assertEqual(len(result), 1)
+        self.assertIs(result[0], group)
+
+    def test_passthrough_for_non_fsdp_config(self) -> None:
+        # SingleDeviceDistributedConfig has no num_sub_groups field; should pass through.
+        group = self._make_group(
+            num_params=4, distributed_config=SingleDeviceDistributedConfig()
+        )
+        result = split_param_groups([group])
+        self.assertEqual(len(result), 1)
+        self.assertIs(result[0], group)
+
+    @parametrize(
+        "num_params, num_sub_groups, expected_sizes",
+        (
+            (6, 3, [2, 2, 2]),
+            (5, 4, [2, 1, 1, 1]),
+        ),
+    )
+    def test_split_distribution(
+        self, num_params: int, num_sub_groups: int, expected_sizes: list[int]
+    ) -> None:
+        config = self._make_fsdp_config(num_sub_groups=num_sub_groups)
+        group = self._make_group(num_params=num_params, distributed_config=config)
+        result = split_param_groups([group])
+        self.assertEqual(len(result), num_sub_groups)
+        self.assertEqual([len(g[PARAMS]) for g in result], expected_sizes)
+
+    def test_preserves_other_group_keys(self) -> None:
+        config = self._make_fsdp_config(num_sub_groups=2)
+        group = self._make_group(
+            num_params=4, distributed_config=config, extra={"lr": 0.01, "beta3": 0.9}
+        )
+        result = split_param_groups([group])
+        for sub_group in result:
+            self.assertEqual(sub_group["lr"], 0.01)
+            self.assertEqual(sub_group["beta3"], 0.9)
+
+    def test_each_sub_group_gets_distinct_config_copy(self) -> None:
+        config = self._make_fsdp_config(num_sub_groups=2)
+        group = self._make_group(num_params=4, distributed_config=config)
+        result = split_param_groups([group])
+        self.assertIsNot(result[0][DISTRIBUTED_CONFIG], result[1][DISTRIBUTED_CONFIG])
+        self.assertIsNot(result[0][DISTRIBUTED_CONFIG], config)
+
+    def test_raises_when_num_sub_groups_exceeds_num_params(self) -> None:
+        config = self._make_fsdp_config(num_sub_groups=5)
+        group = self._make_group(num_params=3, distributed_config=config)
+        with self.assertRaisesRegex(ValueError, "num_sub_groups=5 is too large"):
+            split_param_groups([group])
+
+    def test_splits_each_input_group_independently(self) -> None:
+        config = self._make_fsdp_config(num_sub_groups=2)
+        group_a = self._make_group(num_params=4, distributed_config=config)
+        group_b = self._make_group(
+            num_params=2, distributed_config=SingleDeviceDistributedConfig()
+        )
+        result = split_param_groups([group_a, group_b])
+        # group_a -> 2 sub-groups; group_b passes through unchanged.
+        self.assertEqual(len(result), 3)
+        self.assertEqual(len(result[0][PARAMS]), 2)
+        self.assertEqual(len(result[1][PARAMS]), 2)
+        self.assertIs(result[2], group_b)
+
+    def test_hsdp_raises_when_dominant_param_starves_a_bin(self) -> None:
+        """HSDP: a single huge param can be alone in a bin, leaving fewer
+        than shard_size params there. Greedy bin-packing must detect and reject."""
+        # 6 params split into 2 bins with shard_size=3. Validation passes
+        # (6 // 3 = 2 >= 2). One param has dominant numel; greedy puts it
+        # alone in a bin → 1 < shard_size, violation.
+        params = [torch.zeros(1000)] + [torch.zeros(1) for _ in range(5)]
+        # Avoid real device mesh init (needs distributed setup).
+        mock_mesh = MagicMock()
+        mock_mesh.size.return_value = 3
+        config = HybridShardDistributedConfig(
+            device_mesh=mock_mesh,
+            param_assignment_strategy=FSDPParamAssignmentStrategy.ROUND_ROBIN,
+            num_sub_groups=2,
+        )
+        group: dict[str, object] = {PARAMS: params, DISTRIBUTED_CONFIG: config}
+        with self.assertRaisesRegex(ValueError, "fewer than shard_size=3"):
+            split_param_groups([group])
+
+
+class GreedyBinPackTest(unittest.TestCase):
+    def test_skewed_sizes(self) -> None:
+        """Largest item should be alone; three small items in the other bin."""
+        items = [1000, 100, 100, 100]
+        bins, costs = greedy_bin_pack(items, num_bins=2, cost_fn=lambda x: x)
+        # Bin 0 gets the largest item first, bin 1 gets the rest.
+        self.assertEqual(sorted(costs), [300, 1000])
+        self.assertEqual(sum(len(b) for b in bins), 4)
+
+    def test_more_bins_than_items(self) -> None:
+        """Empty bins should be present when num_bins > len(items)."""
+        items = ["a", "b"]
+        bins, costs = greedy_bin_pack(items, num_bins=4, cost_fn=lambda x: 1)
+        non_empty = [b for b in bins if b]
+        self.assertEqual(len(non_empty), 2)
+        self.assertEqual(len(bins), 4)
+
+    def test_single_bin(self) -> None:
+        """All items should end up in one bin."""
+        items = [10, 20, 30]
+        bins, costs = greedy_bin_pack(items, num_bins=1, cost_fn=lambda x: x)
+        self.assertEqual(len(bins), 1)
+        self.assertEqual(costs, [60])
+        self.assertEqual(sorted(bins[0]), [10, 20, 30])
+
+    def test_equal_sizes(self) -> None:
+        """Equal-cost items should be distributed round-robin across bins."""
+        items = list(range(6))
+        bins, costs = greedy_bin_pack(items, num_bins=3, cost_fn=lambda x: 1)
+        # Each bin should get exactly 2 items.
+        for b in bins:
+            self.assertEqual(len(b), 2)
+        self.assertEqual(costs, [2, 2, 2])
+
+    def test_empty_items(self) -> None:
+        """No items should produce all empty bins."""
+        empty_items: list[int] = []
+        bins, costs = greedy_bin_pack(empty_items, num_bins=3, cost_fn=lambda x: x)
+        self.assertEqual(bins, [[], [], []])
+        self.assertEqual(costs, [0, 0, 0])
diff --git a/gpa/README.md b/gpa/README.md
index 4acba05..ad8c8a6 100644
--- a/gpa/README.md
+++ b/gpa/README.md
@@ -45,20 +45,6 @@ Where:
 - `β₁` and `β₂` are the inner optimizer (AdamW)'s hyperparameters
 - `g^{(t)}` is the gradient coming from the inner optimizer
 
-## Requirements
-
-- PyTorch >= 2.0
-- Python >= 3.10
-- CUDA 11.x or 12.x (for GPU training)
-
-## Installation
-
-The GPAAdamW optimizer is available in the `gpa` package.
-
-```python
-from gpa.gpa_adamw import GPAAdamW
-```
-
 ## Quick Start
 
 ### Basic Usage
@@ -132,7 +118,8 @@ torch.save({
 | `beta1` | EMA coefficient for gradient (like Adam). Must be in [0, 1) | 0.9 |
 | `beta2` | EMA coefficient for squared gradient (like Adam). Must be in [0, 1) | 0.999 |
 | `eps` | Numerical stability term. Must be >= 0 | 1e-8 |
-| `weight_decay` | L2 regularization applied to y-sequence. Must be >= 0 | 0.0 |
+| `weight_decay` | Weight decay coefficient applied as multiplicative shrinkage. Applied to z-sequence by default, or y-sequence if `use_wd_on_y=True`. Must be >= 0 | 0.0 |
+| `use_wd_on_y` | If True, apply weight decay to y-sequence instead of z-sequence | False |
 | `weight_pow_coeff` | Polynomial weighting power, r in the paper (Schedule-Free mode only). Must be >= 0 | 0.0 |
 | `weight_lr_power` | Learning rate weighting power during warmup (Schedule-Free mode only). Must be >= 0 | 2.0 |
 
@@ -247,10 +234,11 @@ optimizer = GPAAdamW(
    | 16                     | 0.9934                  |
    | 32                     | 0.9967                  |
    | 64                     | 0.9984                  |
-   | 128                    | 0.9972                  |
+   | 128                    | 0.9992                  |
 
 4. **Weight Decay:**
-   - Applied to the y-sequence (training sequence)
+   - By default, applied as multiplicative shrinkage on the z-sequence (z-buffer)
+   - Set `use_wd_on_y=True` to apply to the y-sequence (training params) instead
    - Standard values (0.01-0.1) work well
    - Reduce if training becomes unstable
 
@@ -264,61 +252,8 @@ optimizer = GPAAdamW(
 
 GPAAdamW works seamlessly with PyTorch's distributed training APIs including DDP and FSDP. No special configuration is required—simply wrap your model with DDP or FSDP and create the optimizer as usual.
 
-## Checkpointing
-
-### Saving Checkpoints
-
-Save checkpoints in eval mode (optional if y and x sequences are similar):
-
-```python
-
-# Switch to eval mode before saving (optional if y and x sequences are similar)
-model.eval()
-optimizer.eval()
-
-checkpoint = {
-    'epoch': epoch,
-    'model_state_dict': model.state_dict(),
-    'optimizer_state_dict': optimizer.state_dict(),
-}
-torch.save(checkpoint, 'checkpoint.pt')
-```
-
-### Loading Checkpoints
-
-```python
-
-# Create model and optimizer
-model = create_model()
-optimizer = GPAAdamW(model.parameters(), lr=0.001)
-
-# Load checkpoint
-checkpoint = torch.load('checkpoint.pt')
-model.load_state_dict(checkpoint['model_state_dict'])
-optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-start_epoch = checkpoint['epoch']
-
-# Resume training
-model.train()
-optimizer.train()
-
-# ... continue training
-```
-
 ## Running Tests
 
-The test suite is organized into unit tests (`tests/`) and GPU tests (`gpu_tests/`).
-
-### Test Organization
-
-| Test File | Description | # Tests |
-|-----------|-------------|---------|
-| `tests/gpa_adamw_test.py` | Core GPAAdamW unit tests (initialization, step/mode, avg_coeff, state_dict) | 17 |
-| `tests/gpa_equivalence_test.py` | Base optimizer equivalence tests | 1 |
-| `gpu_tests/gpa_adamw_numerics_test.py` | GPU convergence and numerical tests (CPU + CUDA parameterized) | 12 |
-
-### Run Tests
-
 ```bash
 # Core GPAAdamW unit tests
 python -m unittest gpa.tests.gpa_adamw_test -v
@@ -333,55 +268,6 @@ python -m unittest gpa.gpu_tests.gpa_adamw_numerics_test -v
 python -m unittest discover -s gpa.tests -v
 ```
 
-## Common Issues and Troubleshooting
-
-### RuntimeError: Optimizer was not in train mode
-
-**Cause:** Calling `optimizer.step()` without first calling `optimizer.train()`.
-
-**Solution:** Always call `optimizer.train()` before the training loop.
-
-```python
-
-optimizer.train()  # Add this before training
-
-for batch in train_loader:
-    optimizer.zero_grad()
-    # ...
-    optimizer.step()
-```
-
-### Incorrect evaluation results
-
-**Cause:** Evaluating the model without switching to eval mode.
-
-**Solution:** Call `optimizer.eval()` before evaluation (optional if y and x sequences are similar).
-
-```python
-
-model.eval()
-optimizer.eval()  # Add this before evaluation (optional if y ≈ x)
-
-with torch.no_grad():
-    for batch in val_loader:
-        output = model(batch)
-        # ...
-```
-
-### Checkpoint incompatibility
-
-**Cause:** Checkpoint saved in train mode contains y-sequence instead of x-sequence.
-
-**Solution:** Save checkpoints after calling `optimizer.eval()` (optional if y and x sequences are similar).
-
-### Training instability
-
-**Possible solutions:**
-1. Reduce learning rate
-2. Add warmup
-3. Reduce `train_interp_coeff` (e.g., from 0.9 to 0.8)
-4. Increase `beta2` (e.g., from 0.999 to 0.9999)
-
 ## Advanced Options
 
 ### Custom Weighting in Schedule-Free Mode
diff --git a/gpa/__init__.py b/gpa/__init__.py
index e69de29..9071dd2 100644
--- a/gpa/__init__.py
+++ b/gpa/__init__.py
@@ -0,0 +1,8 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+All rights reserved.
+
+This source code is licensed under the BSD-style license found in the
+LICENSE file in the root directory of this source tree.
+
+"""
diff --git a/gpa/gpa_adamw.py b/gpa/gpa_adamw.py
index 4c5b127..fa425f4 100644
--- a/gpa/gpa_adamw.py
+++ b/gpa/gpa_adamw.py
@@ -7,7 +7,6 @@
 
 """
 
-# pyre-unsafe
 from logging import getLogger
 from typing import Callable, Optional, overload, Union
 
@@ -28,13 +27,13 @@
     STEP,
     TRAIN_INTERP_COEFF,
     TRAIN_MODE,
+    USE_WD_ON_Y,
     WEIGHT_DECAY,
     WEIGHT_LR_POWER,
     WEIGHT_POW_COEFF,
     WEIGHT_SUM,
     Z_BUFFER,
 )
-from torch import Tensor
 from torch.optim.optimizer import ParamsT
 
 logger = getLogger()
@@ -102,9 +101,12 @@ class GPAAdamW(torch.optim.Optimizer):
             gradient term. (default: 0.9)
         beta2 (float): Coefficient used for computing running average of the
             gradient squared term. (default: 0.999)
-        weight_decay (float): Weight decay. Note that weight_decay can be either
-            applied to the y or x sequence. In this implementation, it is applied
-            to the y sequence. (default: 0)
+        weight_decay (float): Weight decay coefficient. Applied as multiplicative
+            shrinkage on either the y-sequence or z-sequence depending on
+            use_wd_on_y: z_shrunk = (1 - lr * weight_decay) * z. (default: 0)
+        use_wd_on_y (bool): If True, apply weight decay shrinkage on y (the
+            training params): y *= (1 - lr * weight_decay). If False (default),
+            apply on z (the z-buffer): z *= (1 - lr * weight_decay). (default: False)
         iterate_averaging_type (IterateAveragingType): Controls which averaging
             mode is used. GPA uses a fixed eval_interp_coeff (mu_x).
             SCHEDULE_FREE uses polynomial weighting (no fixed mu_x).
@@ -132,6 +134,7 @@ def __init__(
         beta1: float = 0.9,
         beta2: float = 0.999,
         weight_decay: float = 0,
+        use_wd_on_y: bool = False,
         iterate_averaging_type: IterateAveragingType = IterateAveragingType.GPA,
         train_interp_coeff: float = 0.7,
         eval_interp_coeff: float = 0.9967,
@@ -192,6 +195,7 @@ def __init__(
                 WEIGHT_POW_COEFF: weight_pow_coeff,
                 WEIGHT_LR_POWER: weight_lr_power,
                 WEIGHT_DECAY: weight_decay,
+                USE_WD_ON_Y: use_wd_on_y,
                 EVAL_INTERP_COEFF: eval_interp_coeff,
                 ITERATE_AVERAGING_TYPE: iterate_averaging_type,
             },
@@ -353,8 +357,10 @@ def train(self):
                 self.state[first_param][TRAIN_MODE].fill_(True)
 
     @overload
+    @torch.no_grad()
     def step(self, closure: None = None) -> None: ...
     @overload
+    @torch.no_grad()
     def step(self, closure: Callable[[], float]) -> float: ...
 
     @torch.no_grad()
@@ -386,11 +392,11 @@ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]
             if not group[PARAMS]:
                 continue
 
-            params_with_grad: list[Tensor] = []
-            grads: list[Tensor] = []
-            exp_avgs: list[Tensor] = []
-            exp_avg_sqs: list[Tensor] = []
-            z_buffer_list: list[Tensor] = []
+            params_with_grad: list[torch.Tensor] = []
+            grads: list[torch.Tensor] = []
+            exp_avgs: list[torch.Tensor] = []
+            exp_avg_sqs: list[torch.Tensor] = []
+            z_buffer_list: list[torch.Tensor] = []
 
             self._init_group(
                 group,
@@ -401,8 +407,8 @@ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]
                 z_buffer_list,
             )
 
-            # Get group_first_param for accessing shared state.
-            group_first_param: Tensor = group[PARAMS][0]
+            # Get first_param for accessing shared state.
+            group_first_param: torch.Tensor = group[PARAMS][0]
 
             # Increment step counter and use it as group step.
             self.state[group_first_param][STEP] += 1
@@ -414,6 +420,7 @@ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]
             beta1 = group[BETA1]
             beta2 = group[BETA2]
             weight_decay = group[WEIGHT_DECAY]
+            use_wd_on_y = group[USE_WD_ON_Y]
             weight_lr_power = group[WEIGHT_LR_POWER]
             lr = group[LR]
             weight_pow_coeff = group[WEIGHT_POW_COEFF]
@@ -462,9 +469,12 @@ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]
                 # grad_normalized = exp_avg.div_(denom)
                 grad_normalized = exp_avg.div(bias_correction1).div_(denom)
 
-                # Weight decay calculated at y
+                # Weight decay applied as multiplicative shrinkage.
                 if weight_decay != 0:
-                    grad_normalized.add_(y, alpha=weight_decay)
+                    if use_wd_on_y:
+                        y.mul_(1 - lr * weight_decay)
+                    else:
+                        z.mul_(1 - lr * weight_decay)
 
                 # Memory-efficient y-update without explicitly computing x:
                 # The standard updates are:
diff --git a/gpa/gpa_types.py b/gpa/gpa_types.py
index d1d851a..a6623ca 100644
--- a/gpa/gpa_types.py
+++ b/gpa/gpa_types.py
@@ -25,6 +25,7 @@
 LR_MAX = "lr_max"
 WEIGHT_LR_POWER = "weight_lr_power"
 WEIGHT_DECAY = "weight_decay"
+USE_WD_ON_Y = "use_wd_on_y"
 EVAL_INTERP_COEFF = "eval_interp_coeff"
 ITERATE_AVERAGING_TYPE = "iterate_averaging_type"
 
diff --git a/gpa/tests/gpa_adamw_test.py b/gpa/tests/gpa_adamw_test.py
index 271a88b..e583c7d 100644
--- a/gpa/tests/gpa_adamw_test.py
+++ b/gpa/tests/gpa_adamw_test.py
@@ -725,6 +725,7 @@ def _ref_state_dict(self) -> Dict[str, Any]:
                     "beta1": 0.9,
                     "beta2": 0.999,
                     "weight_decay": 0,
+                    "use_wd_on_y": False,
                     "weight_pow_coeff": 0.0,
                     "weight_lr_power": 2,
                     "eval_interp_coeff": 0.9967,
diff --git a/gpa/tests/gpa_equivalence_test.py b/gpa/tests/gpa_equivalence_test.py
index 7cbb7f8e..c251244 100644
--- a/gpa/tests/gpa_equivalence_test.py
+++ b/gpa/tests/gpa_equivalence_test.py
@@ -19,8 +19,6 @@
     python -m unittest gpa.tests.gpa_equivalence_test -v
 """
 
-# pyre-unsafe
-
 import unittest
 
 import torch
@@ -113,13 +111,78 @@ def test_gpa_with_mu_x_zero_equals_adamw(self) -> None:
 
         Tested with both weight_decay=0 and weight_decay>0. Not bitwise equal
         due to different floating point operation ordering between GPA and
-        PyTorch's AdamW implementation (e.g., GPA applies weight decay
-        by adding it to the search direction while AdamW multiplies it to the parameters directly).
+        PyTorch's AdamW implementation (e.g., GPA applies weight decay as
+        z-shrinkage z.mul_(1 - lr * wd) while AdamW multiplies it to the
+        parameters directly, but the mathematical result is equivalent).
         """
         for weight_decay in [0.0, 0.01]:
             with self.subTest(weight_decay=weight_decay):
                 self._assert_gpa_matches_adamw(weight_decay)
 
+    def test_use_wd_on_y_with_eval_interp_coeff_one(self) -> None:
+        """
+        Test use_wd_on_y behavior when eval_interp_coeff=1.0 (x never updates).
+
+        When eval_interp_coeff=1.0, avg_coeff=0, so the y-lerp is a no-op
+        and z never feeds into y. This means:
+        - use_wd_on_y=False (decay on z) should have no effect on y-parameters,
+          producing the same y as weight_decay=0.
+        - use_wd_on_y=True (decay on y) should produce different y-parameters
+          from weight_decay=0, confirming the decay takes effect.
+
+        This exercises the use_wd_on_y=True code path, which is motivated by
+        LaProp-style weight decay (applied to y=x when train_interp_coeff=1.0).
+        """
+        device = torch.device("cpu")
+        eval_interp_coeff = 1.0
+        weight_decay = 0.01
+
+        y_params = {}
+        for label, wd, use_wd_on_y in [
+            ("no_wd", 0.0, False),
+            ("wd_on_z", weight_decay, False),
+            ("wd_on_y", weight_decay, True),
+        ]:
+            torch.manual_seed(self.seed)
+            model = nn.Linear(10, 5, bias=False).to(device)
+            optimizer = GPAAdamW(
+                model.parameters(),
+                lr=self.lr,
+                beta1=self.beta1,
+                beta2=self.beta2,
+                eps=self.eps,
+                weight_decay=wd,
+                use_wd_on_y=use_wd_on_y,
+                train_interp_coeff=1.0,
+                eval_interp_coeff=eval_interp_coeff,
+                iterate_averaging_type=IterateAveragingType.GPA,
+            )
+
+            torch.manual_seed(100)
+            gradients = [
+                torch.randn(5, 10, device=device) for _ in range(self.num_steps)
+            ]
+
+            optimizer.train()
+            run_optimizer_steps(optimizer, model, gradients)
+
+            y_params[label] = list(model.parameters())[0].data.clone()
+
+        # z-decay should not affect y when eval_interp_coeff=1.0
+        torch.testing.assert_close(
+            y_params["wd_on_z"],
+            y_params["no_wd"],
+            atol=0.0,
+            rtol=0.0,
+            msg="weight decay on z should not affect y when eval_interp_coeff=1.0",
+        )
+
+        # y-decay should affect y
+        self.assertFalse(
+            torch.allclose(y_params["wd_on_y"], y_params["no_wd"]),
+            "weight decay on y should affect y-parameters",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()