Use setter for activation offload policy

Signed-off-by: hongbinl <hongbinl@nvidia.com>

Author: lhb8125

tests/pytorch/test_fusible_ops.py

@@ -101,12 +101,12 @@ def test_basic_operation_activation_offloading_policy(monkeypatch):
     assert calls == [("start", [tensor_id]), ("mark", [tensor_id])]
 
     calls.clear()
-    op.disable_activation_offloading()
+    op.set_activation_offloading(False)
     op.maybe_mark_and_start_activation_offload(tensor, start=True)
     assert calls == [("skip", [tensor_id])]
 
     calls.clear()
-    op.enable_activation_offloading()
+    op.set_activation_offloading(True)
     op.maybe_mark_and_start_activation_offload(tensor, start=True, mark=False)
     assert calls == [("start", [tensor_id])]
 

transformer_engine/pytorch/ops/op.py

@@ -195,17 +195,13 @@ def __init__(self) -> None:
     def is_fused_op(self) -> bool:
         return False
 
-    def disable_activation_offloading(self, disabled: bool = True) -> None:
-        """Disable activation CPU offloading for tensors saved by this op.
+    def set_activation_offloading(self, enabled: bool) -> None:
+        """Enable or disable activation CPU offloading for tensors saved by this op.
 
         CPU offloading is controlled by the surrounding offload context. This setting only
-        opts this operation's saved activation tensors out of that context.
+        opts this operation's saved activation tensors in or out of that context.
         """
-        self.activation_offloading = not disabled
-
-    def enable_activation_offloading(self) -> None:
-        """Re-enable activation CPU offloading for tensors saved by this op."""
-        self.disable_activation_offloading(False)
+        self.activation_offloading = enabled
 
     def maybe_mark_and_start_activation_offload(
         self,