Cache AdamW placement grouping and tensor lists [skip-build]

- DTensor unwrap via _to_local() to eliminate __torch_dispatch__ overhead
- _placement_cache: group params by placement once, reuse every step
- _tensor_cache: cache (params_local, moment1, moment2, state_steps)
- Only gradients collected fresh each step
- Profiler annotations for AdamW stages

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

torch-ext/optimizer/adamw.py

@@ -1,8 +1,12 @@
+import logging
 from collections import defaultdict
 from typing import cast
 
 import torch
 from torch.distributed.tensor import DTensor
+from torch.profiler import record_function
+
+logger = logging.getLogger(__name__)
 
 
 def fused_adamw(
@@ -72,54 +76,72 @@ def fused_adamw(
         )
 
 
-def step_adamw_params(optimizer_state, params, group):
-    """Run fused AdamW on a list of parameters sharing the same placement.
+def _to_local(t):
+    """Unwrap DTensor to local tensor for fused ops."""
+    return t._local_tensor if isinstance(t, DTensor) else t
 
-    Args:
-        optimizer_state: The optimizer's state dict (self.state in Muon).
-        params: List of parameters to update.
-        group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay.
-    """
+
+# ---------------------------------------------------------------------------
+# Caches for eliminating per-step Python overhead.
+#
+# Placement grouping and tensor list assembly are identical every step
+# (params don't change placement, moment/step tensors are the same objects
+# after initialisation).  We cache them keyed by id() of the param list
+# stored in param_groups (stable across steps).
+#
+# Only gradients change each step and must be collected fresh.
+# ---------------------------------------------------------------------------
+
+# id(group["params"]) → dict[placement_key, list[param]]
+_placement_cache: dict[int, dict[tuple, list]] = {}
+
+# id(placement_group_list) → (params_local, moment1, moment2, state_steps)
+_tensor_cache: dict[int, tuple[list, list, list, list]] = {}
+
+
+def _step_adamw_params_slow(optimizer_state, params, group):
+    """Uncached fallback for the rare case where some params lack grads."""
     params_with_grads = []
     grads = []
     moment1 = []
     moment2 = []
-    max_exp_avg_sqs = []
     state_steps = []
-    lr = group["lr"]
-    beta1, beta2 = group["adamw_betas"]
-    eps = group["adamw_eps"]
-    weight_decay = group["weight_decay"]
 
     for p in params:
         g = p.grad
         if g is None:
             continue
         state = optimizer_state[p]
-        params_with_grads.append(p)
-        grads.append(g)
+        params_with_grads.append(_to_local(p))
+        grads.append(_to_local(g))
         if "step" not in state:
-            state["step"] = (torch.zeros((),
-                                         dtype=torch.float32,
-                                         device=p.device))
+            state["step"] = torch.zeros((),
+                                        dtype=torch.float32,
+                                        device=p.device)
             state["moment1"] = torch.zeros_like(g)
             state["moment2"] = torch.zeros_like(g)
-        moment1.append(state["moment1"])
-        moment2.append(state["moment2"])
+        moment1.append(_to_local(state["moment1"]))
+        moment2.append(_to_local(state["moment2"]))
         if not isinstance(state["step"], torch.Tensor):
-            step_tensor = torch.tensor(state["step"],
-                                       dtype=torch.float32,
-                                       device=p.device)
-        else:
-            step_tensor = state["step"]
-        state_steps.append(step_tensor)
+            state["step"] = torch.tensor(state["step"],
+                                         dtype=torch.float32,
+                                         device=p.device)
+        state_steps.append(state["step"])
+
+    if not params_with_grads:
+        return
+
+    lr = group["lr"]
+    beta1, beta2 = group["adamw_betas"]
+    eps = group["adamw_eps"]
+    weight_decay = group["weight_decay"]
 
     fused_adamw(
         params_with_grads,
         grads,
         moment1,
         moment2,
-        max_exp_avg_sqs,
+        [],
         state_steps,
         amsgrad=False,
         beta1=beta1,
@@ -131,24 +153,119 @@ def step_adamw_params(optimizer_state, params, group):
     )
 
 
+def step_adamw_params(optimizer_state, params, group):
+    """Run fused AdamW on a list of parameters sharing the same placement.
+
+    After the first call, cached tensor lists (params_local, moment1,
+    moment2, state_steps) are reused — only gradients are collected fresh.
+
+    Args:
+        optimizer_state: The optimizer's state dict (self.state in Muon).
+        params: List of parameters to update.
+        group: Parameter group dict with lr, adamw_betas, adamw_eps, weight_decay.
+    """
+    # Collect grads — the only thing that changes each step.
+    with record_function("adamw::collect_grads"):
+        grads = []
+        for p in params:
+            g = p.grad
+            if g is None:
+                # Rare: fall back to slow path that filters per-param.
+                _step_adamw_params_slow(optimizer_state, params, group)
+                return
+            grads.append(_to_local(g))
+
+    tensor_key = id(params)
+    if tensor_key not in _tensor_cache:
+        with record_function("adamw::init_tensor_cache"):
+            params_local = []
+            moment1 = []
+            moment2 = []
+            state_steps = []
+
+            for p in params:
+                state = optimizer_state[p]
+                params_local.append(_to_local(p))
+                if "step" not in state:
+                    state["step"] = torch.zeros((),
+                                                dtype=torch.float32,
+                                                device=p.device)
+                    state["moment1"] = torch.zeros_like(p.grad)
+                    state["moment2"] = torch.zeros_like(p.grad)
+                moment1.append(_to_local(state["moment1"]))
+                moment2.append(_to_local(state["moment2"]))
+                if not isinstance(state["step"], torch.Tensor):
+                    state["step"] = torch.tensor(state["step"],
+                                                 dtype=torch.float32,
+                                                 device=p.device)
+                state_steps.append(state["step"])
+
+            _tensor_cache[tensor_key] = (params_local, moment1, moment2,
+                                         state_steps)
+
+    params_local, moment1, moment2, state_steps = _tensor_cache[tensor_key]
+
+    lr = group["lr"]
+    beta1, beta2 = group["adamw_betas"]
+    eps = group["adamw_eps"]
+    weight_decay = group["weight_decay"]
+
+    with record_function("adamw::fused_adamw"):
+        fused_adamw(
+            params_local,
+            grads,
+            moment1,
+            moment2,
+            [],
+            state_steps,
+            amsgrad=False,
+            beta1=beta1,
+            beta2=beta2,
+            lr=lr,
+            weight_decay=weight_decay,
+            eps=eps,
+            maximize=False,
+        )
+
+
 def step_adamw(optimizer_state, group):
     """Dispatch AdamW step, grouping parameters by type and placement.
 
+    Placement grouping is cached after the first call since params never
+    change their placement between steps.
+
     Args:
         optimizer_state: The optimizer's state dict (self.state in Muon).
         group: Parameter group dict.
     """
     params = group["params"]
+    placement_key = id(params)
 
-    # group params with its type and placement
-    placement_to_params: dict[tuple, list[torch.Tensor]] = defaultdict(list)
-    for p in params:
-        match p:
-            case DTensor():
-                placement_to_params[tuple([p.placements,
-                                           p.device_mesh])].append(p)
-            case torch.Tensor():
-                placement_to_params[tuple([torch.Tensor, None])].append(p)
-
-    for group_params in placement_to_params.values():
+    if placement_key not in _placement_cache:
+        with record_function("adamw::group_by_placement"):
+            placement_to_params: dict[tuple,
+                                      list[torch.Tensor]] = defaultdict(list)
+            for p in params:
+                match p:
+                    case DTensor():
+                        logger.debug(
+                            "[AdamW] DTensor param: shape=%s, placements=%s, "
+                            "mesh=%s, grad=%s", p.shape, p.placements,
+                            p.device_mesh.mesh_dim_names,
+                            p.grad.shape if p.grad is not None else None)
+                        placement_to_params[tuple(
+                            [p.placements, p.device_mesh])].append(p)
+                    case torch.Tensor():
+                        logger.debug(
+                            "[AdamW] plain param: shape=%s, grad=%s", p.shape,
+                            p.grad.shape if p.grad is not None else None)
+                        placement_to_params[tuple([torch.Tensor,
+                                                   None])].append(p)
+
+            logger.debug("[AdamW] %d placement groups, %d total params",
+                         len(placement_to_params), len(params))
+
+            _placement_cache[placement_key] = dict(placement_to_params)
+
+    for group_params in _placement_cache[placement_key].values():
         step_adamw_params(optimizer_state, group_params, group)