fix(checkpointing): harden precision state CPU staging

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/training/optim.py

@@ -139,6 +139,23 @@ def _get_te_fused_adam_class() -> type[torch.optim.Optimizer] | None:
     return cast(type[torch.optim.Optimizer], FusedAdam)
 
 
+def _cpu_staging_get_unscaled_state(
+    fallback: Callable[..., object],
+) -> Callable[..., torch.Tensor]:
+    """Wrap a TE state accessor without depending on its call signature."""
+
+    def _get_unscaled_state_on_cpu(*args: object, **kwargs: object) -> torch.Tensor:
+        state = fallback(*args, **kwargs)
+        if not isinstance(state, torch.Tensor):
+            raise TypeError(
+                "Transformer Engine FusedAdam.get_unscaled_state() must return a torch.Tensor "
+                f"for CPU checkpoint staging, but returned {type(state).__name__}."
+            )
+        return state.cpu()
+
+    return _get_unscaled_state_on_cpu
+
+
 @contextmanager
 def memory_efficient_precision_aware_optimizer_state_checkpointing(
     optimizer: MegatronOptimizer | None,
@@ -196,21 +213,15 @@ def memory_efficient_precision_aware_optimizer_state_checkpointing(
             ):
                 continue
 
-            original_get_unscaled_state: Callable[..., torch.Tensor] = inner.get_unscaled_state
-
-            def _get_unscaled_state_on_cpu(
-                fused_adam: torch.optim.Optimizer,
-                param: torch.nn.Parameter,
-                state_name: str,
-                skip_unscale: bool = False,
-                *,
-                _fallback: Callable[..., torch.Tensor] = original_get_unscaled_state,
-            ) -> torch.Tensor:
-                del fused_adam
-                return _fallback(param, state_name, skip_unscale).cpu()
+            original_get_unscaled_state = getattr(inner, "get_unscaled_state", None)
+            if not callable(original_get_unscaled_state):
+                raise RuntimeError(
+                    "CPU checkpoint staging requires Transformer Engine FusedAdam.get_unscaled_state() "
+                    "to be callable. The installed Transformer Engine checkpoint API is incompatible."
+                )
 
             previous_instance_method = inner.__dict__.get("get_unscaled_state", missing_method)
-            setattr(inner, "get_unscaled_state", MethodType(_get_unscaled_state_on_cpu, inner))
+            setattr(inner, "get_unscaled_state", _cpu_staging_get_unscaled_state(original_get_unscaled_state))
             patched.append((inner, previous_instance_method))
 
         if patched:

tests/unit_tests/training/test_optim.py

@@ -194,10 +194,12 @@ def get_unscaled_state(
         param: torch.nn.Parameter,
         state_name: str,
         skip_unscale: bool = False,
+        *,
+        multiplier: float = 1.0,
     ) -> torch.Tensor:
         del skip_unscale
         self.get_unscaled_state_calls += 1
-        return self.state[param][state_name].float()
+        return self.state[param][state_name].float() * multiplier
 
 
 class _FakeParamRange:
@@ -275,6 +277,43 @@ def test_stages_unscaled_state_on_cpu_and_restores_method(self):
 
         assert inner.get_unscaled_state_calls == 2
 
+    def test_forwards_positional_and_keyword_arguments(self):
+        """The wrapper stays compatible when TE extends its accessor signature."""
+        distributed, inner, param = self._distributed_optimizer()
+
+        with patch("megatron.bridge.training.optim._get_te_fused_adam_class", return_value=_FakeFusedAdam):
+            with memory_efficient_precision_aware_optimizer_state_checkpointing(distributed, enabled=True):
+                state = inner.get_unscaled_state(param, "exp_avg", True, multiplier=3.0)
+
+        torch.testing.assert_close(state, torch.full((4,), 3.0))
+
+    def test_rejects_non_tensor_state_and_restores_instance_method(self):
+        """TE return-contract drift fails clearly without leaking the patch."""
+        distributed, inner, param = self._distributed_optimizer()
+        original_instance_method = MagicMock(return_value="not a tensor")
+        inner.get_unscaled_state = original_instance_method
+
+        with (
+            patch("megatron.bridge.training.optim._get_te_fused_adam_class", return_value=_FakeFusedAdam),
+            pytest.raises(TypeError, match="must return a torch.Tensor"),
+        ):
+            with memory_efficient_precision_aware_optimizer_state_checkpointing(distributed, enabled=True):
+                inner.get_unscaled_state(param, "exp_avg")
+
+        assert inner.__dict__["get_unscaled_state"] is original_instance_method
+
+    def test_rejects_missing_te_state_accessor(self):
+        """An incompatible TE API fails before checkpoint construction begins."""
+        distributed, inner, _ = self._distributed_optimizer()
+        inner.get_unscaled_state = None
+
+        with (
+            patch("megatron.bridge.training.optim._get_te_fused_adam_class", return_value=_FakeFusedAdam),
+            pytest.raises(RuntimeError, match=r"FusedAdam\.get_unscaled_state\(\).*callable"),
+        ):
+            with memory_efficient_precision_aware_optimizer_state_checkpointing(distributed, enabled=True):
+                pass
+
     @pytest.mark.parametrize("incompatibility", ["fp32", "cpu_offload", "fsdp", "stub"])
     def test_does_not_patch_incompatible_optimizer(self, incompatibility: str):
         state_dtype = torch.float32 if incompatibility == "fp32" else torch.bfloat16
@@ -336,6 +375,59 @@ def test_restores_method_when_checkpointing_raises(self):
 
         assert "get_unscaled_state" not in inner.__dict__
 
+    @pytest.mark.run_only_on("gpu")
+    def test_real_te_fused_adam_stages_state_and_restores_method(self):
+        """The pinned TE precision-aware optimizer returns CPU checkpoint state."""
+        te_optimizers = pytest.importorskip("transformer_engine.pytorch.optimizers")
+        fused_adam_class = te_optimizers.FusedAdam
+        param = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16, device="cuda"))
+        inner = fused_adam_class(
+            [param],
+            master_weights=True,
+            master_weight_dtype=torch.float16,
+            exp_avg_dtype=torch.bfloat16,
+            exp_avg_sq_dtype=torch.bfloat16,
+            use_decoupled_grad=True,
+        )
+        inner.initialize_state(param, store_param_remainders=False)
+        distributed = _FakeDistribOpt(model_param=param, shard_main_param=param, inner=inner)
+        distributed.config.use_precision_aware_optimizer = True
+
+        assert "get_unscaled_state" not in inner.__dict__
+        with memory_efficient_precision_aware_optimizer_state_checkpointing(distributed, enabled=True) as patched:
+            portable_state = next(iter(inner.state_dict()["state"].values()))
+            assert patched == 1
+            assert portable_state
+            assert all(state.device.type == "cpu" for state in portable_state.values())
+
+        assert "get_unscaled_state" not in inner.__dict__
+        assert inner.get_unscaled_state(param, "exp_avg").device.type == "cuda"
+
+    @pytest.mark.run_only_on("gpu")
+    def test_real_te_fused_adam_all_fp32_state_is_not_patched(self):
+        """The real TE optimizer keeps its native path when no expansion is needed."""
+        te_optimizers = pytest.importorskip("transformer_engine.pytorch.optimizers")
+        fused_adam_class = te_optimizers.FusedAdam
+        param = torch.nn.Parameter(torch.zeros(4, dtype=torch.bfloat16, device="cuda"))
+        inner = fused_adam_class(
+            [param],
+            master_weights=True,
+            master_weight_dtype=torch.float32,
+            exp_avg_dtype=torch.float32,
+            exp_avg_sq_dtype=torch.float32,
+            use_decoupled_grad=True,
+        )
+        inner.initialize_state(param, store_param_remainders=False)
+        distributed = _FakeDistribOpt(model_param=param, shard_main_param=param, inner=inner)
+        distributed.config.use_precision_aware_optimizer = True
+
+        with memory_efficient_precision_aware_optimizer_state_checkpointing(distributed, enabled=True) as patched:
+            native_state = next(iter(inner.state_dict()["state"].values()))
+            assert patched == 0
+            assert all(state.device.type == "cuda" for state in native_state.values())
+
+        assert "get_unscaled_state" not in inner.__dict__
+
 
 class TestMemoryEfficientFp32OptimizerStateLoading:
     """Tests for the scoped TE FusedAdam checkpoint-load fast path."""