fix(megatron): isolate FP32 adapter optimizer state

skyrl/backends/skyrl_train/workers/megatron/adapter_store.py

@@ -26,6 +26,17 @@ def iter_opts(opt) -> List[Any]:
     return [opt]
 
 
+def _iter_main_param_groups(opt) -> Iterable[List[torch.Tensor]]:
+    """Yield both mixed-precision masters and native FP32 optimizer shards.
+
+    Native FP32 shards alias DDP parameter buffers, but their optimizer state
+    still needs to travel with the adapter when the live buffers are swapped.
+    Keep the same group order for slot allocation, snapshot, and restore.
+    """
+    yield from getattr(opt, "shard_fp32_from_float16_groups", None) or []
+    yield from getattr(opt, "shard_fp32_groups", None) or []
+
+
 def _iter_buffers(model_chunks) -> Iterable[Tuple[int, int, Any]]:
     """Yield (mc_idx, buf_idx, buffer) for every LoRA-trainable DDP buffer."""
     for mc_idx, mc in enumerate(model_chunks):
@@ -124,7 +135,8 @@ class AdapterSlot:
           forward_backward aren't lost when another tenant runs before
           this adapter's optim_step.
       cpu_main_param[opt_idx][g] -> list[Tensor], shapes matching
-          opt.shard_fp32_from_float16_groups[g].
+          groups in opt.shard_fp32_from_float16_groups followed by
+          opt.shard_fp32_groups.
       cpu_opt_state[opt_idx][g][i] -> dict[str, Tensor], mirroring
           opt.optimizer.state[main_param] for every tensor-valued entry
           (exp_avg, exp_avg_sq, step, ...).
@@ -196,7 +208,7 @@ def _is_scalar_state(v: Any) -> bool:
 
     def _allocate_empty_slot(self, model_chunks, optimizer) -> AdapterSlot:
         slot = AdapterSlot()
-        # Param data + grad data: one pinned bf16 tensor each per (mc, buffer).
+        # Param data + grad data: one pinned tensor each per (mc, buffer).
         # Grads must travel with the slot so forward_backward calls accumulate
         # into the correct adapter even when requests from tenants interleave.
         for mc_idx, _buf_idx, buf in _iter_buffers(model_chunks):
@@ -209,8 +221,7 @@ def _allocate_empty_slot(self, model_chunks, optimizer) -> AdapterSlot:
         for _opt in iter_opts(optimizer):
             opt_main: List[List[torch.Tensor]] = []
             opt_state: List[List[dict]] = []
-            groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or []
-            for g, group in enumerate(groups):
+            for group in _iter_main_param_groups(_opt):
                 main_g: List[torch.Tensor] = []
                 state_g: List[dict] = []
                 for main_param in group:
@@ -263,8 +274,7 @@ def _snapshot(self, slot: AdapterSlot, model_chunks, optimizer) -> None:
             if _is_resident(buf.grad_data):
                 slot.cpu_grad_data[mc_idx][buf_idx].copy_(buf.grad_data, non_blocking=True)
         for opt_idx, _opt in enumerate(iter_opts(optimizer)):
-            groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or []
-            for g, group in enumerate(groups):
+            for g, group in enumerate(_iter_main_param_groups(_opt)):
                 for i, main_param in enumerate(group):
                     slot.cpu_main_param[opt_idx][g][i].copy_(main_param, non_blocking=True)
                     state = _opt.optimizer.state.get(main_param, {})
@@ -294,8 +304,7 @@ def _restore(self, slot: AdapterSlot, model_chunks, optimizer) -> None:
             if _is_resident(buf.grad_data):
                 buf.grad_data.copy_(slot.cpu_grad_data[mc_idx][buf_idx], non_blocking=True)
         for opt_idx, _opt in enumerate(iter_opts(optimizer)):
-            groups = getattr(_opt, "shard_fp32_from_float16_groups", None) or []
-            for g, group in enumerate(groups):
+            for g, group in enumerate(_iter_main_param_groups(_opt)):
                 for i, main_param in enumerate(group):
                     main_param.copy_(slot.cpu_main_param[opt_idx][g][i], non_blocking=True)
                     state = _opt.optimizer.state.get(main_param, {})

tests/backends/skyrl_train/workers/test_adapter_store_optimizer_state.py

@@ -0,0 +1,166 @@
+"""CPU numerical tests for per-adapter DistributedOptimizer state.
+
+Only Megatron's CUDA-dependent container types and CPU pinning are replaced.
+The production AdapterStore runs all allocation/copy/swap operations, and real
+PyTorch Adam updates are compared with separately trained adapter baselines.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+
+import pytest
+import torch
+
+
+class _DDP:
+    def __init__(self, params):
+        self.params = params
+        self.buffers = [
+            SimpleNamespace(param_data=p.detach(), grad_data=torch.zeros_like(p), params=[p]) for p in params
+        ]
+        self.expert_parallel_buffers = []
+
+    def named_parameters(self):
+        return [(f"adapter.weight_{i}", p) for i, p in enumerate(self.params)]
+
+
+class _ChainedOptimizer:
+    def __init__(self, optimizers):
+        self.chained_optimizers = optimizers
+
+
+@pytest.fixture
+def adapter_store_module(monkeypatch):
+    # Load the real source in an isolated module namespace so CUDA-free tests
+    # neither require Megatron nor leave fake Megatron imports in other tests.
+    modules = {
+        "megatron": {},
+        "megatron.core": {},
+        "megatron.core.parallel_state": {"get_data_parallel_group": lambda: None},
+        "megatron.core.distributed": {"DistributedDataParallel": _DDP},
+        "megatron.core.optimizer": {"ChainedOptimizer": _ChainedOptimizer},
+    }
+    for name, attrs in modules.items():
+        module = ModuleType(name)
+        module.__path__ = []
+        module.__dict__.update(attrs)
+        monkeypatch.setitem(sys.modules, name, module)
+    path = Path(__file__).resolve().parents[4] / "skyrl/backends/skyrl_train/workers/megatron/adapter_store.py"
+    spec = importlib.util.spec_from_file_location("_adapter_store_cpu_test", path)
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, spec.name, module)
+    spec.loader.exec_module(module)
+    monkeypatch.setattr(module, "_new_pinned_like", lambda t: torch.empty_like(t, device="cpu"))
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(synchronize=lambda: None))
+    monkeypatch.setattr(module.dist, "is_initialized", lambda: False)
+    return module
+
+
+class _TrainingState:
+    """DDP/shard layout with real Adam state and rank-local parameter views."""
+
+    def __init__(self, layouts, group_step):
+        self.group_step = group_step
+        self.chunks = []
+        self.optimizers = []
+        self.param_pairs = []
+        for dtypes in layouts:
+            models = [torch.nn.Parameter(torch.tensor([1.0, 2.0], dtype=dtype)) for dtype in dtypes]
+            # Only the second element belongs to this rank's optimizer shard.
+            # FP32 shards share DDP storage; BF16 parameters have FP32 masters.
+            mains = [
+                torch.nn.Parameter(model.detach()[1:] if model.dtype == torch.float32 else model.detach()[1:].float())
+                for model in models
+            ]
+            adam = torch.optim.Adam([{"params": []}, {"params": mains}], lr=0.01, betas=(0.9, 0.999))
+            for main in mains:
+                adam.state[main] = {"exp_avg": torch.zeros_like(main), "exp_avg_sq": torch.zeros_like(main)}
+                if not group_step:
+                    adam.state[main]["step"] = torch.tensor(0.0)
+            if group_step:
+                for group in adam.param_groups:
+                    group["step"] = 0
+            # Preserve empty groups, as on a rank without shards of a dtype.
+            opt = SimpleNamespace(
+                optimizer=adam,
+                shard_fp32_from_float16_groups=[
+                    [],
+                    [main for model, main in zip(models, mains) if model.dtype == torch.bfloat16],
+                ],
+                shard_fp32_groups=[[], [main for model, main in zip(models, mains) if model.dtype == torch.float32]],
+            )
+            self.optimizers.append(opt)
+            self.chunks.append(_DDP(models))
+            self.param_pairs.extend(zip(models, mains))
+        self.optimizer = self.optimizers[0] if len(self.optimizers) == 1 else _ChainedOptimizer(self.optimizers)
+
+    def step(self, gradient):
+        for index, (_, main) in enumerate(self.param_pairs):
+            main.grad = torch.full_like(main, gradient * (index + 1))
+        for opt in self.optimizers:
+            adam = opt.optimizer
+            if self.group_step:
+                # TE FusedAdam stores the counter on the group, whereas
+                # PyTorch Adam consumes a per-parameter counter. Translate
+                # that storage convention just around the real Adam update.
+                for group in adam.param_groups:
+                    for main in group["params"]:
+                        adam.state[main]["step"] = torch.tensor(float(group["step"]))
+            adam.step()
+            adam.zero_grad(set_to_none=True)
+            if self.group_step:
+                for group in adam.param_groups:
+                    group["step"] += 1
+                    for main in group["params"]:
+                        del adam.state[main]["step"]
+        with torch.no_grad():
+            for model, main in self.param_pairs:
+                model[1:].copy_(main)
+
+    def assert_matches(self, expected):
+        for (model, main), (expected_model, expected_main) in zip(self.param_pairs, expected.param_pairs):
+            torch.testing.assert_close(model, expected_model, rtol=0, atol=0)
+            torch.testing.assert_close(main, expected_main, rtol=0, atol=0)
+        for opt, expected_opt in zip(self.optimizers, expected.optimizers):
+            actual_state = opt.optimizer.state_dict()
+            expected_state = expected_opt.optimizer.state_dict()
+            assert actual_state["param_groups"] == expected_state["param_groups"]
+            torch.testing.assert_close(actual_state["state"], expected_state["state"], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("group_step", [False, True], ids=["pytorch-counter", "te-group-counter"])
+@pytest.mark.parametrize(
+    "layouts",
+    [
+        [(torch.float32,)],
+        [(torch.bfloat16,)],
+        [(torch.bfloat16, torch.float32)],
+        [(torch.bfloat16,), (torch.float32,)],
+    ],
+    ids=["fp32", "bf16", "mixed", "chained"],
+)
+def test_interleaved_adapters_match_independent_adam(adapter_store_module, layouts, group_step):
+    live = _TrainingState(layouts, group_step)
+    baselines = {name: _TrainingState(layouts, group_step) for name in ("a", "b")}
+    store = adapter_store_module.AdapterStore()
+    signature = object()
+    store.register_pristine(live.chunks, live.optimizer, signature)
+    store.create("a", live.chunks, live.optimizer, signature)
+
+    # Train A before B is created so the fresh slot must come from pristine,
+    # including zero moments/counters, rather than from A's live optimizer.
+    for _ in range(20):
+        live.step(1.0)
+        baselines["a"].step(1.0)
+    live.assert_matches(baselines["a"])
+    store.create("b", live.chunks, live.optimizer, signature)
+
+    for name, gradient in [("b", -1.0), ("a", 0.3), ("b", -0.7), ("a", -0.2)]:
+        store.swap_to(name, live.chunks, live.optimizer)
+        live.step(gradient)
+        baselines[name].step(gradient)
+        # A leaked positive momentum makes B's first negative-gradient update
+        # move in the wrong direction, even when its group counter is restored.
+        live.assert_matches(baselines[name])