feat: CPUOffloadPool.reload_group에 sync_streams 파라미터 추가 (#31)

* feat: tag-based per-group reload for CPU offload + manual offload mode

Extend CPUOffloadPool + Muon so callers can drive offload/reload
explicitly (layer-lockstep overlap) rather than implicitly inside
optimizer.step.

CPUOffloadPool (torch-ext/optimizer/cpu_offload.py):
* ``track(tensor, tag=None)`` — optional tag per managed tensor.
* ``reload_group(tag, sync_streams=())`` — reload just the tensors
tagged with ``tag``; the reload stream ``wait_stream`` s on each
entry in ``sync_streams`` before issuing H2D. This avoids
allocator cross-stream reuse races under
``PYTORCH_ALLOC_CONF=expandable_segments:True``: if the block
returned by ``storage.resize_`` was last touched on an FSDP
all-gather stream, waiting on that stream enforces FIFO ordering
between the prior use and our H2D write.
* ``reload_untagged()`` — bulk-reload everything not attached to a
tag (for the non-expert portion of the optimizer state in layer-
lockstep flows).
* ``wait_reload()`` is now self-clearing (resets ``_reload_event``
after one wait).

Muon (torch-ext/optimizer/muon.py):
* ``manual_offload`` flag: when set, ``step()`` skips its own
``reload`` / ``offload`` calls so the caller can drive them.
* ``set_param_tags(id(param) -> tag)``: propagated to the pool on
first offload so ``reload_group`` picks the right tensors up.
* New public helpers: ``reload_group``, ``reload_untagged``,
``wait_reload``, ``offload`` — mirroring pool semantics.
* Baseline (non-manual) path: add explicit ``wait_reload`` after
``pool.reload()`` in ``step`` and ``turn_off_cpu_offload``.

Default ``sync_streams=()`` keeps existing callers behaviourally
unchanged.

* style: yapf

* chore: whitelist math notation 'Ot' for typos hook

_typos.toml

@@ -0,0 +1,3 @@
+[default.extend-words]
+# Math notation used in docs/muon-clip.md (O subscript t, update step output)
+Ot = "Ot"

test/test_cpu_offload.py

@@ -29,7 +29,8 @@ def _setup():
 
 
 def _make_mesh(world_size):
-    return dist.init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp",))
+    return dist.init_device_mesh("cuda", (world_size, ),
+                                 mesh_dim_names=("dp", ))
 
 
 def test_correctness(rank, world_size):
@@ -47,11 +48,12 @@ def test_correctness(rank, world_size):
     num_steps = 3
 
     # Pre-generate all data on all ranks (same seed → same values).
-    full_params = [torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)]
-    full_grads = [
-        [torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)]
-        for _ in range(num_steps)
+    full_params = [
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)
     ]
+    full_grads = [[
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)
+    ] for _ in range(num_steps)]
 
     def make_optimizer(cpu_offload):
         params, names = [], []
@@ -60,19 +62,17 @@ def test_correctness(rank, world_size):
             p = torch.nn.Parameter(dt)
             params.append(p)
             names.append(f"layer.{i}.weight")
-        param_groups = [
-            {
-                "params": params,
-                "names": names,
-                "use_muon": True,
-                "lr": 0.02,
-                "weight_decay": 0.01,
-                "momentum": 0.95,
-                "nesterov": True,
-                "ns_steps": 5,
-                "none_grad": False,
-            }
-        ]
+        param_groups = [{
+            "params": params,
+            "names": names,
+            "use_muon": True,
+            "lr": 0.02,
+            "weight_decay": 0.01,
+            "momentum": 0.95,
+            "nesterov": True,
+            "ns_steps": 5,
+            "none_grad": False,
+        }]
         optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
         if cpu_offload:
             optim.turn_on_cpu_offload()
@@ -121,25 +121,22 @@ def test_memory(rank, world_size):
         full = torch.randn(dim0, dim1, device="cuda")
         dt = distribute_tensor(full, mesh, [Shard(0)])
         p = torch.nn.Parameter(dt)
-        p.grad = distribute_tensor(
-            torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-        )
+        p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                   mesh, [Shard(0)])
         params.append(p)
         names.append(f"layer.{i}.weight")
 
-    param_groups = [
-        {
-            "params": params,
-            "names": names,
-            "use_muon": True,
-            "lr": 0.02,
-            "weight_decay": 0.01,
-            "momentum": 0.95,
-            "nesterov": True,
-            "ns_steps": 5,
-            "none_grad": False,
-        }
-    ]
+    param_groups = [{
+        "params": params,
+        "names": names,
+        "use_muon": True,
+        "lr": 0.02,
+        "weight_decay": 0.01,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": 5,
+        "none_grad": False,
+    }]
     optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
     optim.turn_on_cpu_offload()
 
@@ -155,8 +152,7 @@ def test_memory(rank, world_size):
         local_buf = buf._local_tensor if isinstance(buf, DTensor) else buf
         assert local_buf.untyped_storage().size() == 0, (
             f"Expected freed GPU storage after offload, got "
-            f"{local_buf.untyped_storage().size()} bytes"
-        )
+            f"{local_buf.untyped_storage().size()} bytes")
 
     # Verify CPU pool has pinned buffers.
     pool = optim._cpu_offload_pool
@@ -166,9 +162,8 @@ def test_memory(rank, world_size):
 
     # Run another step to verify reload + compute + offload cycle works.
     for p in params:
-        p.grad = distribute_tensor(
-            torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-        )
+        p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                   mesh, [Shard(0)])
     optim.step()
     torch.cuda.synchronize()
 
@@ -217,26 +212,21 @@ def test_adamw_offload(rank, world_size):
         adamw_names.append(f"layer.{i}.bias")
 
     # Pre-generate grads.
-    muon_grads = [
-        [torch.randn(64, 128, device="cuda") for _ in range(4)]
-        for _ in range(num_steps)
-    ]
-    adamw_grads = [
-        [torch.randn(128, device="cuda") for _ in range(3)] for _ in range(num_steps)
-    ]
+    muon_grads = [[torch.randn(64, 128, device="cuda") for _ in range(4)]
+                  for _ in range(num_steps)]
+    adamw_grads = [[torch.randn(128, device="cuda") for _ in range(3)]
+                   for _ in range(num_steps)]
 
     def make_optimizer(cpu_offload):
         mp = [
             torch.nn.Parameter(
-                distribute_tensor(p.data.full_tensor().clone(), mesh, [Shard(0)])
-            )
-            for p in muon_params
+                distribute_tensor(p.data.full_tensor().clone(), mesh,
+                                  [Shard(0)])) for p in muon_params
         ]
         ap = [
             torch.nn.Parameter(
-                distribute_tensor(p.data.full_tensor().clone(), mesh, [Shard(0)])
-            )
-            for p in adamw_params
+                distribute_tensor(p.data.full_tensor().clone(), mesh,
+                                  [Shard(0)])) for p in adamw_params
         ]
         param_groups = [
             {
@@ -306,8 +296,7 @@ def test_adamw_offload(rank, world_size):
             t = state[key]
             local_t = t._local_tensor if isinstance(t, DTensor) else t
             assert local_t.untyped_storage().size() == 0, (
-                f"AdamW {key} storage not freed after offload"
-            )
+                f"AdamW {key} storage not freed after offload")
 
     set_ns_compile(True)
     if rank == 0:
@@ -335,25 +324,22 @@ def test_memory_savings(rank, world_size):
             full = torch.randn(dim0, dim1, device="cuda")
             dt = distribute_tensor(full, mesh, [Shard(0)])
             p = torch.nn.Parameter(dt)
-            p.grad = distribute_tensor(
-                torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-            )
+            p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                       mesh, [Shard(0)])
             params.append(p)
             names.append(f"layer.{i}.weight")
 
-        param_groups = [
-            {
-                "params": params,
-                "names": names,
-                "use_muon": True,
-                "lr": 0.02,
-                "weight_decay": 0.01,
-                "momentum": 0.95,
-                "nesterov": True,
-                "ns_steps": 5,
-                "none_grad": False,
-            }
-        ]
+        param_groups = [{
+            "params": params,
+            "names": names,
+            "use_muon": True,
+            "lr": 0.02,
+            "weight_decay": 0.01,
+            "momentum": 0.95,
+            "nesterov": True,
+            "ns_steps": 5,
+            "none_grad": False,
+        }]
         optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
         if cpu_offload:
             optim.turn_on_cpu_offload()
@@ -370,16 +356,17 @@ def test_memory_savings(rank, world_size):
     mem_with_offload = run_step(True)
 
     if rank == 0:
-        logger.info("Memory without offload: %.2f MB", mem_no_offload / 1024**2)
-        logger.info("Memory with offload:    %.2f MB", mem_with_offload / 1024**2)
+        logger.info("Memory without offload: %.2f MB",
+                    mem_no_offload / 1024**2)
+        logger.info("Memory with offload:    %.2f MB",
+                    mem_with_offload / 1024**2)
         saved = mem_no_offload - mem_with_offload
         logger.info("Memory saved:           %.2f MB", saved / 1024**2)
 
     assert mem_with_offload < mem_no_offload, (
         f"Expected memory reduction with CPU offload. "
         f"Without: {mem_no_offload / 1024**2:.2f} MB, "
-        f"With: {mem_with_offload / 1024**2:.2f} MB"
-    )
+        f"With: {mem_with_offload / 1024**2:.2f} MB")
 
     set_ns_compile(True)
     if rank == 0:
@@ -400,11 +387,12 @@ def test_toggle_correctness(rank, world_size):
     num_params = 4
     num_steps = 6
 
-    full_params = [torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)]
-    full_grads = [
-        [torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)]
-        for _ in range(num_steps)
+    full_params = [
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)
     ]
+    full_grads = [[
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_params)
+    ] for _ in range(num_steps)]
 
     def make_optimizer():
         params, names = [], []
@@ -413,19 +401,17 @@ def test_toggle_correctness(rank, world_size):
             p = torch.nn.Parameter(dt)
             params.append(p)
             names.append(f"layer.{i}.weight")
-        param_groups = [
-            {
-                "params": params,
-                "names": names,
-                "use_muon": True,
-                "lr": 0.02,
-                "weight_decay": 0.01,
-                "momentum": 0.95,
-                "nesterov": True,
-                "ns_steps": 5,
-                "none_grad": False,
-            }
-        ]
+        param_groups = [{
+            "params": params,
+            "names": names,
+            "use_muon": True,
+            "lr": 0.02,
+            "weight_decay": 0.01,
+            "momentum": 0.95,
+            "nesterov": True,
+            "ns_steps": 5,
+            "none_grad": False,
+        }]
         optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
         return optim, params
 
@@ -446,7 +432,8 @@ def test_toggle_correctness(rank, world_size):
         for i in range(num_params):
             g = full_grads[step_idx][i]
             params_ref[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)])
-            params_toggle[i].grad = distribute_tensor(g.clone(), mesh, [Shard(0)])
+            params_toggle[i].grad = distribute_tensor(g.clone(), mesh,
+                                                      [Shard(0)])
 
         optim_ref.step()
         optim_toggle.step()
@@ -492,19 +479,17 @@ def test_leak(rank, world_size):
         params.append(p)
         names.append(f"layer.{i}.weight")
 
-    param_groups = [
-        {
-            "params": params,
-            "names": names,
-            "use_muon": True,
-            "lr": 0.02,
-            "weight_decay": 0.01,
-            "momentum": 0.95,
-            "nesterov": True,
-            "ns_steps": 5,
-            "none_grad": False,
-        }
-    ]
+    param_groups = [{
+        "params": params,
+        "names": names,
+        "use_muon": True,
+        "lr": 0.02,
+        "weight_decay": 0.01,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": 5,
+        "none_grad": False,
+    }]
     optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
     optim.turn_on_cpu_offload()
 
@@ -519,9 +504,8 @@ def test_leak(rank, world_size):
 
     for step_idx in range(num_steps):
         for p in params:
-            p.grad = distribute_tensor(
-                torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-            )
+            p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                       mesh, [Shard(0)])
 
         optim.step()
         torch.cuda.synchronize()
@@ -564,8 +548,7 @@ def test_leak(rank, world_size):
     # GPU memory should not grow beyond warmup baseline.
     assert gpu_final <= gpu_after_warmup, (
         f"GPU memory leak detected! Warmup: {gpu_after_warmup / 1024**2:.2f} MB, "
-        f"Final: {gpu_final / 1024**2:.2f} MB"
-    )
+        f"Final: {gpu_final / 1024**2:.2f} MB")
 
     # CPU RSS should not grow more than 50 MB over warmup (allows for minor
     # Python/CUDA runtime overhead but catches real leaks).
@@ -573,12 +556,12 @@ def test_leak(rank, world_size):
     assert cpu_growth < 50, (
         f"CPU memory leak detected! Growth: {cpu_growth:.2f} MB over "
         f"{num_steps - 2} steps (warmup={cpu_after_warmup:.2f} MB, "
-        f"final={cpu_final:.2f} MB)"
-    )
+        f"final={cpu_final:.2f} MB)")
 
     set_ns_compile(True)
     if rank == 0:
-        logger.info("PASSED: test_leak (GPU stable, CPU growth=%.2f MB)", cpu_growth)
+        logger.info("PASSED: test_leak (GPU stable, CPU growth=%.2f MB)",
+                    cpu_growth)
 
 
 def test_state_dict_save_load(rank, world_size):
@@ -606,28 +589,26 @@ def test_state_dict_save_load(rank, world_size):
     num_steps = 3
 
     # Pre-generate all data.
-    muon_init = [torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon)]
-    adamw_init = [torch.randn(dim1, device="cuda") for _ in range(num_adamw)]
-    all_grads_muon = [
-        [torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon)]
-        for _ in range(num_steps * 2)
-    ]
-    all_grads_adamw = [
-        [torch.randn(dim1, device="cuda") for _ in range(num_adamw)]
-        for _ in range(num_steps * 2)
+    muon_init = [
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon)
     ]
+    adamw_init = [torch.randn(dim1, device="cuda") for _ in range(num_adamw)]
+    all_grads_muon = [[
+        torch.randn(dim0, dim1, device="cuda") for _ in range(num_muon)
+    ] for _ in range(num_steps * 2)]
+    all_grads_adamw = [[
+        torch.randn(dim1, device="cuda") for _ in range(num_adamw)
+    ] for _ in range(num_steps * 2)]
 
     def make_optimizer(cpu_offload):
         mp = [
             torch.nn.Parameter(
-                distribute_tensor(muon_init[i].clone(), mesh, [Shard(0)])
-            )
+                distribute_tensor(muon_init[i].clone(), mesh, [Shard(0)]))
             for i in range(num_muon)
         ]
         ap = [
             torch.nn.Parameter(
-                distribute_tensor(adamw_init[i].clone(), mesh, [Shard(0)])
-            )
+                distribute_tensor(adamw_init[i].clone(), mesh, [Shard(0)]))
             for i in range(num_adamw)
         ]
         param_groups = [
@@ -666,17 +647,15 @@ def test_state_dict_save_load(rank, world_size):
     for step_idx in range(num_steps):
         for i in range(num_muon):
             mp_off[i].grad = distribute_tensor(
-                all_grads_muon[step_idx][i].clone(), mesh, [Shard(0)]
-            )
+                all_grads_muon[step_idx][i].clone(), mesh, [Shard(0)])
         for i in range(num_adamw):
             ap_off[i].grad = distribute_tensor(
-                all_grads_adamw[step_idx][i].clone(), mesh, [Shard(0)]
-            )
+                all_grads_adamw[step_idx][i].clone(), mesh, [Shard(0)])
         optim_off.step()
 
     with pytest.raises(
-        RuntimeError, match="turn_off_cpu_offload\\(\\) before checkpoint save"
-    ):
+            RuntimeError,
+            match="turn_off_cpu_offload\\(\\) before checkpoint save"):
         optim_off.state_dict()
 
     optim_off.turn_off_cpu_offload()
@@ -688,8 +667,7 @@ def test_state_dict_save_load(rank, world_size):
             if isinstance(val, torch.Tensor) and val.is_floating_point():
                 assert val.untyped_storage().size() > 0, (
                     f"state_dict() returned empty storage for key '{key}' — "
-                    f"offload reload is broken"
-                )
+                    f"offload reload is broken")
 
     if rank == 0:
         logger.info("state_dict() contains valid (non-empty) tensors")
@@ -724,8 +702,8 @@ def test_state_dict_save_load(rank, world_size):
         for i in range(num_adamw):
             ap_ref[i].data = ap_off[i].data.clone()
         with pytest.raises(
-            RuntimeError, match="turn_off_cpu_offload\\(\\) before checkpoint load"
-        ):
+                RuntimeError,
+                match="turn_off_cpu_offload\\(\\) before checkpoint load"):
             optim_ref.load_state_dict(copy.deepcopy(sd_off))
         optim_ref.turn_off_cpu_offload()
         optim_ref.load_state_dict(copy.deepcopy(sd_off))
@@ -749,8 +727,8 @@ def test_state_dict_save_load(rank, world_size):
                 if flat_key in flat_target:
                     param_state[key] = flat_target[flat_key]
         with pytest.raises(
-            RuntimeError, match="turn_off_cpu_offload\\(\\) before checkpoint load"
-        ):
+                RuntimeError,
+                match="turn_off_cpu_offload\\(\\) before checkpoint load"):
             optim_resumed.load_state_dict(copy.deepcopy(sd_loaded))
         optim_resumed.turn_off_cpu_offload()
         optim_resumed.load_state_dict(sd_loaded)
@@ -795,8 +773,7 @@ def test_state_dict_save_load(rank, world_size):
             buf = state["momentum_buffer"]
             local_buf = buf._local_tensor if isinstance(buf, DTensor) else buf
             assert local_buf.untyped_storage().size() == 0, (
-                "Resumed optimizer should have offloaded state after step()"
-            )
+                "Resumed optimizer should have offloaded state after step()")
 
     set_ns_compile(True)
     if rank == 0:
@@ -821,25 +798,22 @@ def test_checkpoint_memory(rank, world_size):
         full = torch.randn(dim0, dim1, device="cuda")
         dt = distribute_tensor(full, mesh, [Shard(0)])
         p = torch.nn.Parameter(dt)
-        p.grad = distribute_tensor(
-            torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-        )
+        p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                   mesh, [Shard(0)])
         params.append(p)
         names.append(f"layer.{i}.weight")
 
-    param_groups = [
-        {
-            "params": params,
-            "names": names,
-            "use_muon": True,
-            "lr": 0.02,
-            "weight_decay": 0.01,
-            "momentum": 0.95,
-            "nesterov": True,
-            "ns_steps": 5,
-            "none_grad": False,
-        }
-    ]
+    param_groups = [{
+        "params": params,
+        "names": names,
+        "use_muon": True,
+        "lr": 0.02,
+        "weight_decay": 0.01,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": 5,
+        "none_grad": False,
+    }]
     optim = Muon(params=param_groups, chunk_size=2, warmup_step=1)
     optim.turn_on_cpu_offload()
 
@@ -867,8 +841,8 @@ def test_checkpoint_memory(rank, world_size):
         )
 
     with pytest.raises(
-        RuntimeError, match="turn_off_cpu_offload\\(\\) before checkpoint save"
-    ):
+            RuntimeError,
+            match="turn_off_cpu_offload\\(\\) before checkpoint save"):
         optim.state_dict()
 
     optim.turn_off_cpu_offload()
@@ -885,28 +859,24 @@ def test_checkpoint_memory(rank, world_size):
     assert mem_after_turn_off > mem_after_step, (
         f"turn_off_cpu_offload() should reload states to GPU. "
         f"After offload: {mem_after_step / 1024**2:.2f} MB, "
-        f"After turn_off: {mem_after_turn_off / 1024**2:.2f} MB"
-    )
+        f"After turn_off: {mem_after_turn_off / 1024**2:.2f} MB")
 
     optim.turn_on_cpu_offload()
     torch.cuda.synchronize()
     mem_after_turn_on = torch.cuda.memory_allocated()
 
     if rank == 0:
-        logger.info(
-            "After turn_on_cpu_offload: GPU alloc=%.2f MB", mem_after_turn_on / 1024**2
-        )
+        logger.info("After turn_on_cpu_offload: GPU alloc=%.2f MB",
+                    mem_after_turn_on / 1024**2)
 
     assert mem_after_turn_on <= mem_after_step + 4 * 1024 * 1024, (
         f"turn_on_cpu_offload() should return memory to offloaded level. "
         f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), "
-        f"got {mem_after_turn_on / 1024**2:.2f} MB"
-    )
+        f"got {mem_after_turn_on / 1024**2:.2f} MB")
 
     for p in params:
-        p.grad = distribute_tensor(
-            torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-        )
+        p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                   mesh, [Shard(0)])
     optim.step()
     torch.cuda.synchronize()
 
@@ -922,12 +892,11 @@ def test_checkpoint_memory(rank, world_size):
     assert mem_after_next_step <= mem_after_step + 4 * 1024 * 1024, (
         f"Memory should return to offloaded level after step(). "
         f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), "
-        f"got {mem_after_next_step / 1024**2:.2f} MB"
-    )
+        f"got {mem_after_next_step / 1024**2:.2f} MB")
 
     with pytest.raises(
-        RuntimeError, match="turn_off_cpu_offload\\(\\) before checkpoint load"
-    ):
+            RuntimeError,
+            match="turn_off_cpu_offload\\(\\) before checkpoint load"):
         optim.load_state_dict(copy.deepcopy(sd_for_load))
 
     optim.turn_off_cpu_offload()
@@ -943,24 +912,21 @@ def test_checkpoint_memory(rank, world_size):
         )
 
     assert mem_after_load >= mem_after_turn_off, (
-        "Loaded optimizer state should stay on GPU while offload is disabled"
-    )
+        "Loaded optimizer state should stay on GPU while offload is disabled")
 
     optim.turn_on_cpu_offload()
     torch.cuda.synchronize()
 
     pool = optim._cpu_offload_pool
     assert pool._initialized, (
-        "Offload pool should be initialized after re-enabling offload"
-    )
+        "Offload pool should be initialized after re-enabling offload")
     for grp in pool._groups.values():
         assert grp["cpu_flat"].is_pinned(), "CPU buffer must be pinned"
 
     # Step 5: verify the loaded optimizer can still step correctly.
     for p in params:
-        p.grad = distribute_tensor(
-            torch.randn(dim0, dim1, device="cuda"), mesh, [Shard(0)]
-        )
+        p.grad = distribute_tensor(torch.randn(dim0, dim1, device="cuda"),
+                                   mesh, [Shard(0)])
     optim.step()
     torch.cuda.synchronize()
 
@@ -968,8 +934,7 @@ def test_checkpoint_memory(rank, world_size):
     assert mem_final <= mem_after_step + 4 * 1024 * 1024, (
         f"Final memory should be at offloaded level. "
         f"Expected <= {mem_after_step / 1024**2:.2f} MB (+4 MB tolerance), "
-        f"got {mem_final / 1024**2:.2f} MB"
-    )
+        f"got {mem_final / 1024**2:.2f} MB")
 
     set_ns_compile(True)
     if rank == 0:

torch-ext/optimizer/cpu_offload.py

@@ -20,6 +20,7 @@ from collections import defaultdict
 
 import torch
 from torch.distributed.tensor import DTensor
+from torch.profiler import record_function
 
 logger = logging.getLogger(__name__)
 
@@ -35,6 +36,9 @@ class CPUOffloadPool:
     def __init__(self):
         self._managed: list[torch.Tensor] = []
         self._storage_nbytes: dict[int, int] = {}  # id(t) → bytes
+        # Optional tag → managed-indices map for group-wise reload
+        # (e.g. per-layer lockstep reload driven by backward hooks).
+        self._tag_to_indices: dict[str, list[int]] = {}
 
         # Per-dtype group: populated on first offload.
         # dtype → dict with keys:
@@ -45,6 +49,8 @@ class CPUOffloadPool:
         self._groups: dict[torch.dtype, dict] = {}
 
         self._offload_stream: torch.cuda.Stream | None = None
+        self._reload_stream: torch.cuda.Stream | None = None
+        self._reload_event: torch.cuda.Event | None = None
         self._device: torch.device | None = None
         self._initialized: bool = False
         self._logged: bool = False
@@ -59,9 +65,28 @@ class CPUOffloadPool:
         if self._offload_stream is None:
             self._offload_stream = torch.cuda.Stream(device=self._device)
 
+    def _ensure_reload_stream(self):
+        if self._reload_stream is None:
+            least_priority, _ = torch.cuda.Stream.priority_range()
+            self._reload_stream = torch.cuda.Stream(
+                device=self._device,
+                priority=least_priority,
+            )
+            logger.info(
+                "[CPUOffload] reload stream created with priority=%d "
+                "(range: %d..%d)",
+                least_priority,
+                *torch.cuda.Stream.priority_range(),
+            )
+
     # ------------------------------------------------------------------
-    def track(self, tensor: torch.Tensor):
-        """Register a GPU tensor for CPU offloading.  Idempotent."""
+    def track(self, tensor: torch.Tensor, tag: str | None = None):
+        """Register a GPU tensor for CPU offloading.  Idempotent.
+
+        If ``tag`` is given, the tensor's managed index is recorded under
+        that tag so callers can trigger a partial reload via
+        :meth:`reload_group`.
+        """
         tid = id(tensor)
         if tid in self._storage_nbytes:
             return
@@ -73,7 +98,10 @@ class CPUOffloadPool:
         if storage.size() == 0:
             return
         self._storage_nbytes[tid] = storage.size()
+        idx = len(self._managed)
         self._managed.append(tensor)
+        if tag is not None:
+            self._tag_to_indices.setdefault(tag, []).append(idx)
 
     # ------------------------------------------------------------------
     def _init_buffers(self):
@@ -93,7 +121,10 @@ class CPUOffloadPool:
                 indices.append(idx)
                 offsets.append((off, n))
                 off += n
-            cpu_flat = torch.empty(off, dtype=dtype, device="cpu", pin_memory=True)
+            cpu_flat = torch.empty(off,
+                                   dtype=dtype,
+                                   device="cpu",
+                                   pin_memory=True)
             self._groups[dtype] = {
                 "indices": indices,
                 "offsets": offsets,
@@ -137,7 +168,8 @@ class CPUOffloadPool:
                 for i, mgd_idx in enumerate(indices):
                     local = self._local(self._managed[mgd_idx])
                     off, n = offsets[i]
-                    cpu_flat[off : off + n].copy_(local.reshape(-1), non_blocking=True)
+                    cpu_flat[off:off + n].copy_(local.reshape(-1),
+                                                non_blocking=True)
 
                 offloaded_bytes += grp["total"] * cpu_flat.element_size()
 
@@ -151,8 +183,7 @@ class CPUOffloadPool:
                 raise RuntimeError(
                     f"Tensor storage is already freed (size=0) before offload. "
                     f"This indicates a double-free or external interference. "
-                    f"Tensor shape: {t.shape}, dtype: {t.dtype}"
-                )
+                    f"Tensor shape: {t.shape}, dtype: {t.dtype}")
 
         if not self._logged:
             logger.info(
@@ -162,45 +193,172 @@ class CPUOffloadPool:
 
     # ------------------------------------------------------------------
     def reload(self):
-        """Per-tensor H2D from CPU flat buffer on the default stream.
+        """Per-tensor H2D from CPU flat buffer.
+
+        Storage re-allocation (``resize_``) runs on the current (default)
+        stream.  H2D copies run on a dedicated ``_reload_stream``.
 
-        Runs on the current (default) CUDA stream to avoid stream
-        interaction issues with the parallel Muon pipeline.  Since
-        pinned CPU memory is the source, the copies overlap with
-        GPU idle time between steps.
+        Call :meth:`wait_reload` before consuming the reloaded tensors.
         """
         if not self._managed or not self._initialized:
             return
+        self._ensure_reload_stream()
 
         reloaded_bytes = 0
 
-        # Re-allocate all GPU storages first.
-        for t in self._managed:
-            local = self._local(t)
-            storage = local.untyped_storage()
-            if storage.size() != 0:
-                raise RuntimeError(
-                    f"Storage should have been freed (size=0) before reload, "
-                    f"but got size={storage.size()}. "
-                    f"Tensor shape: {t.shape}, dtype: {t.dtype}"
-                )
-            storage.resize_(self._storage_nbytes[id(t)])
-
-        # Per-tensor H2D copies from CPU flat buffer slices.
-        # non_blocking=True with pinned source allows DMA overlap.
-        for dtype, grp in self._groups.items():
-            indices = grp["indices"]
-            offsets = grp["offsets"]
-            cpu_flat = grp["cpu_flat"]
-
-            for i, mgd_idx in enumerate(indices):
-                local = self._local(self._managed[mgd_idx])
-                off, n = offsets[i]
-                local.reshape(-1).copy_(cpu_flat[off : off + n], non_blocking=True)
-
-            reloaded_bytes += grp["total"] * cpu_flat.element_size()
+        # Re-allocate all GPU storages with per-tensor profiling.
+        with record_function("CPUOffload::resize_storages"):
+            for i, t in enumerate(self._managed):
+                local = self._local(t)
+                storage = local.untyped_storage()
+                if storage.size() != 0:
+                    raise RuntimeError(
+                        f"Storage should have been freed (size=0) before reload, "
+                        f"but got size={storage.size()}. "
+                        f"Tensor shape: {t.shape}, dtype: {t.dtype}")
+                nbytes = self._storage_nbytes[id(t)]
+                with record_function(f"resize_[{i}]_{nbytes // 1024}KB"):
+                    storage.resize_(nbytes)
+
+        # Reload stream waits for the resize_ ops to finish.
+        alloc_event = torch.cuda.current_stream(self._device).record_event()
+        self._reload_stream.wait_event(alloc_event)
+
+        # Per-tensor H2D copies on the reload stream.
+        with record_function("CPUOffload::h2d_copies"):
+            with torch.cuda.stream(self._reload_stream):
+                for dtype, grp in self._groups.items():
+                    indices = grp["indices"]
+                    offsets = grp["offsets"]
+                    cpu_flat = grp["cpu_flat"]
+
+                    for i, mgd_idx in enumerate(indices):
+                        local = self._local(self._managed[mgd_idx])
+                        off, n = offsets[i]
+                        local.reshape(-1).copy_(cpu_flat[off:off + n],
+                                                non_blocking=True)
+
+                    reloaded_bytes += grp["total"] * cpu_flat.element_size()
+
+        self._reload_event = self._reload_stream.record_event()
 
         if not self._logged:
             logger.info(
-                "[CPUOffload] Reloaded %.2f MB (CPU → GPU)", reloaded_bytes / (1024**2)
+                "[CPUOffload] Reloaded %.2f MB (CPU → GPU, async)",
+                reloaded_bytes / (1024**2),
             )
+            self._logged = True
+
+    def reload_group(self, tag: str, sync_streams: tuple = ()):
+        """Reload only the managed tensors registered under ``tag``.
+
+        Intended for layer-lockstep overlap: backward frees a layer's
+        activations, then the backward hook calls ``reload_group`` with
+        that layer's tag so the H2D copy reuses the freshly-freed memory
+        from the default stream's allocator pool.
+
+        ``sync_streams`` is an optional iterable of CUDA streams whose
+        currently-queued work must complete before the H2D memcpy runs.
+        This is used to avoid allocator cross-stream reuse races under
+        ``expandable_segments``: if a just-freed block was last used on
+        FSDP's all-gather stream, making the reload stream wait on that
+        stream guarantees FIFO ordering between the block's prior use
+        and our H2D write.
+        """
+        if not self._managed or not self._initialized:
+            return
+        indices = self._tag_to_indices.get(tag)
+        if not indices:
+            return
+        self._ensure_reload_stream()
+
+        # Sync reload_stream with the supplied streams (e.g. FSDP AG
+        # streams) before we queue any H2D: ensures past uses of any
+        # allocator block we're about to reuse are fully drained.
+        for s in sync_streams:
+            if s is not None:
+                self._reload_stream.wait_stream(s)
+
+        idx_set = set(indices)
+
+        with record_function(f"CPUOffload::group_resize[{tag}]"):
+            for i in indices:
+                t = self._managed[i]
+                local = self._local(t)
+                storage = local.untyped_storage()
+                if storage.size() == 0:
+                    storage.resize_(self._storage_nbytes[id(t)])
+
+        alloc_event = torch.cuda.current_stream(self._device).record_event()
+        self._reload_stream.wait_event(alloc_event)
+
+        with record_function(f"CPUOffload::group_h2d[{tag}]"):
+            with torch.cuda.stream(self._reload_stream):
+                for dtype, grp in self._groups.items():
+                    indices_grp = grp["indices"]
+                    offsets = grp["offsets"]
+                    cpu_flat = grp["cpu_flat"]
+
+                    for i, mgd_idx in enumerate(indices_grp):
+                        if mgd_idx not in idx_set:
+                            continue
+                        local = self._local(self._managed[mgd_idx])
+                        off, n = offsets[i]
+                        local.reshape(-1).copy_(cpu_flat[off:off + n],
+                                                non_blocking=True)
+
+        self._reload_event = self._reload_stream.record_event()
+
+    def reload_untagged(self):
+        """Reload managed tensors that were not registered under any tag.
+
+        Useful when a subset of params (e.g. MoE experts) is driven via
+        per-tag layer-lockstep hooks while the remainder should still be
+        reloaded before optimizer.step() in a single bulk call.
+        """
+        if not self._managed or not self._initialized:
+            return
+        tagged: set[int] = set()
+        for idx_list in self._tag_to_indices.values():
+            tagged.update(idx_list)
+        untagged = [i for i in range(len(self._managed)) if i not in tagged]
+        if not untagged:
+            return
+        self._ensure_reload_stream()
+
+        idx_set = set(untagged)
+
+        with record_function("CPUOffload::untagged_resize"):
+            for i in untagged:
+                t = self._managed[i]
+                local = self._local(t)
+                storage = local.untyped_storage()
+                if storage.size() == 0:
+                    storage.resize_(self._storage_nbytes[id(t)])
+
+        alloc_event = torch.cuda.current_stream(self._device).record_event()
+        self._reload_stream.wait_event(alloc_event)
+
+        with record_function("CPUOffload::untagged_h2d"):
+            with torch.cuda.stream(self._reload_stream):
+                for dtype, grp in self._groups.items():
+                    indices_grp = grp["indices"]
+                    offsets = grp["offsets"]
+                    cpu_flat = grp["cpu_flat"]
+
+                    for i, mgd_idx in enumerate(indices_grp):
+                        if mgd_idx not in idx_set:
+                            continue
+                        local = self._local(self._managed[mgd_idx])
+                        off, n = offsets[i]
+                        local.reshape(-1).copy_(cpu_flat[off:off + n],
+                                                non_blocking=True)
+
+        self._reload_event = self._reload_stream.record_event()
+
+    def wait_reload(self):
+        """Block the current (default) stream until reload H2D completes."""
+        if self._reload_event is not None:
+            torch.cuda.current_stream(self._device).wait_event(
+                self._reload_event)
+            self._reload_event = None

torch-ext/optimizer/muon.py

@@ -242,8 +242,12 @@ class Muon(torch.optim.Optimizer):
         self.use_distributed_muon = use_distributed_muon
         self.expert_keys = expert_keys
         self.cpu_offload = False
+        self.manual_offload = False
         self._cpu_offload_pool: CPUOffloadPool | None = None
         self._offload_initialized = False
+        # id(param) -> tag, consumed by _register_states_for_offload so the
+        # offload pool can do group-wise reload (e.g. per-layer lockstep).
+        self._param_tags: dict[int, str] = {}
         self._parallel_cache: dict[tuple[str, ...], dict] = {}
         self._expert_expand_cache: dict[tuple[int, ...], dict] = {}
 
@@ -955,15 +959,16 @@ class Muon(torch.optim.Optimizer):
                 if p not in self.state:
                     continue
                 state = self.state[p]
+                tag = self._param_tags.get(id(p))
                 if group.get("use_muon", False):
                     if "momentum_buffer" in state:
-                        pool.track(state["momentum_buffer"])
+                        pool.track(state["momentum_buffer"], tag=tag)
                         tracked += 1
                 else:
                     if "moment1" in state:
-                        pool.track(state["moment1"])
+                        pool.track(state["moment1"], tag=tag)
                     if "moment2" in state:
-                        pool.track(state["moment2"])
+                        pool.track(state["moment2"], tag=tag)
                         tracked += 1
         logger.info("[CPUOffload] Registered %d param states for offload",
                     tracked)
@@ -986,8 +991,10 @@ class Muon(torch.optim.Optimizer):
                 loss = closure()
 
         # H2D: reload optimizer states from CPU before computation.
-        if self.cpu_offload and self._offload_initialized:
-            self._cpu_offload_pool.reload()
+        if not self.manual_offload:
+            if self.cpu_offload and self._offload_initialized:
+                self._cpu_offload_pool.reload()
+                self._cpu_offload_pool.wait_reload()
 
         logger.debug("[Muon.step] expert_keys=%s, %d param groups",
                      self.expert_keys, len(self.param_groups))
@@ -1004,6 +1011,53 @@ class Muon(torch.optim.Optimizer):
                 step_adamw(self.state, group)
 
         # D2H: offload optimizer states to CPU after computation.
+        if not self.manual_offload:
+            if self.cpu_offload:
+                if not self._offload_initialized:
+                    if self._cpu_offload_pool is None:
+                        self._cpu_offload_pool = CPUOffloadPool()
+                    self._register_states_for_offload()
+                    self._offload_initialized = True
+                self._cpu_offload_pool.offload()
+
+        return loss
+
+    # ------------------------------------------------------------------
+    # CPU offload public helpers
+    # ------------------------------------------------------------------
+
+    def reload_group(self, tag: str, sync_streams: tuple = ()):
+        """Reload optimizer states registered under ``tag``.
+
+        Tags are set via :meth:`set_param_tags` before the first step.
+        ``sync_streams`` forwards to :meth:`CPUOffloadPool.reload_group`
+        so callers (e.g. FSDP pre/post-hook patches) can make the reload
+        stream wait on collective streams before its H2D runs.
+        """
+        if self.cpu_offload and self._offload_initialized:
+            self._cpu_offload_pool.reload_group(tag, sync_streams=sync_streams)
+
+    def reload_untagged(self):
+        """Reload all optimizer states not attached to any tag."""
+        if self.cpu_offload and self._offload_initialized:
+            self._cpu_offload_pool.reload_untagged()
+
+    def set_param_tags(self, param_tags: dict[int, str]) -> None:
+        """Attach an ``id(param) -> tag`` mapping for group-wise reload.
+
+        Must be called before the first ``step()`` (i.e. before
+        :meth:`_register_states_for_offload`) so the pool receives tags
+        when states are first registered.
+        """
+        self._param_tags = dict(param_tags)
+
+    def wait_reload(self):
+        """Block the default stream until the async reload completes."""
+        if self.cpu_offload and self._offload_initialized:
+            self._cpu_offload_pool.wait_reload()
+
+    def offload(self):
+        """Offload optimizer states from GPU to CPU (D2H)."""
         if self.cpu_offload:
             if not self._offload_initialized:
                 if self._cpu_offload_pool is None:
@@ -1012,12 +1066,6 @@ class Muon(torch.optim.Optimizer):
                 self._offload_initialized = True
             self._cpu_offload_pool.offload()
 
-        return loss
-
-    # ------------------------------------------------------------------
-    # CPU offload public helpers
-    # ------------------------------------------------------------------
-
     def turn_on_cpu_offload(self):
         """Enable CPU offload for optimizer states."""
         if self.cpu_offload:
@@ -1039,6 +1087,7 @@ class Muon(torch.optim.Optimizer):
         logger.info("[Muon] turn_off_cpu_offload")
         if self._offload_initialized:
             self._cpu_offload_pool.reload()
+            self._cpu_offload_pool.wait_reload()
             torch.cuda.current_stream().synchronize()
         self._cpu_offload_pool = None
         self._offload_initialized = False

torch-ext/optimizer/newton_schulz.py

@@ -32,30 +32,28 @@ def _optimal_quintic(l, u, max_iter=1000):
     E = inf
     for _ in range(max_iter):
         old_E = E
-        LHS = np.array(
-            [
-                [l, l**3, l**5, 1],
-                [q, q**3, q**5, -1],
-                [r, r**3, r**5, 1],
-                [u, u**3, u**5, -1],
-            ]
-        )
+        LHS = np.array([
+            [l, l**3, l**5, 1],
+            [q, q**3, q**5, -1],
+            [r, r**3, r**5, 1],
+            [u, u**3, u**5, -1],
+        ])
         a, b, c, E = np.linalg.solve(LHS, np.ones(4))
         if not np.all(np.isfinite([a, b, c, E])):
             raise ValueError(
                 f"_optimal_quintic: non-finite solve result a={a}, b={b}, c={c}, E={E}"
             )
         q, r = np.sqrt(
-            (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) / (10 * c)
-        )
+            (-3 * b + np.array([-1, 1]) * sqrt(9 * b**2 - 20 * a * c)) /
+            (10 * c))
         if not np.all(np.isfinite([q, r])):
-            raise ValueError(f"_optimal_quintic: non-finite node update q={q}, r={r}")
+            raise ValueError(
+                f"_optimal_quintic: non-finite node update q={q}, r={r}")
         if abs(old_E - E) <= 1e-15:
             break
     else:
         raise RuntimeError(
-            f"_optimal_quintic: did not converge after {max_iter} iterations"
-        )
+            f"_optimal_quintic: did not converge after {max_iter} iterations")
     return float(a), float(b), float(c)
 
 
@@ -114,9 +112,10 @@ def _optimal_composition(l, num_iters, safety_factor_eps=0, cushion=0):
 #   - Polar Express: analytically optimal per step, adapting to the shrinking
 #     singular-value interval [l, u] as iterations progress; converges all
 #     singular values to 1, producing the exact polar factor UV^T.
-_coeffs_list = _optimal_composition(
-    l=1e-3, num_iters=10, safety_factor_eps=1e-2, cushion=0.02
-)
+_coeffs_list = _optimal_composition(l=1e-3,
+                                    num_iters=10,
+                                    safety_factor_eps=1e-2,
+                                    cushion=0.02)
 
 
 # This code is adapted from:
@@ -150,8 +149,7 @@ def _zeropower_via_newtonschulz5(G, steps):
 
     X = X / (X.norm() + 1e-7)
     hs = _coeffs_list[:steps] + list(
-        repeat(_coeffs_list[-1], steps - len(_coeffs_list))
-    )
+        repeat(_coeffs_list[-1], steps - len(_coeffs_list)))
     buf1 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     buf2 = torch.empty(X.size(0), X.size(0), dtype=X.dtype, device=X.device)
     # Perform the NS iterations
@@ -186,8 +184,7 @@ def _zeropower_via_newtonschulz5_batched(G, steps):
     X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
 
     hs = _coeffs_list[:steps] + list(
-        repeat(_coeffs_list[-1], steps - len(_coeffs_list))
-    )
+        repeat(_coeffs_list[-1], steps - len(_coeffs_list)))
     for a, b, c in hs:
         buf1 = torch.bmm(X, X.transpose(-2, -1))
         buf2 = torch.bmm(buf1, buf1.transpose(-2, -1))