Drop TE cached quantized weights before offloading the training actor

miles/backends/megatron_utils/actor.py

@@ -23,7 +23,7 @@
 from miles.utils.distributed_utils import get_gloo_group
 from miles.utils.ft_utils.indep_dp import IndepDPInfo
 from miles.utils.hf_config import load_hf_config
-from miles.utils.memory_utils import clear_memory, print_memory
+from miles.utils.memory_utils import clear_memory, clear_quantized_weight_workspaces, print_memory
 from miles.utils.multi_lora import is_multi_lora_enabled
 from miles.utils.processing_utils import load_tokenizer
 from miles.utils.ray_utils import Box
@@ -302,6 +302,10 @@ def sleep(self) -> None:
             logger.info("sleep() called while already offloaded; skipping")
             return
 
+        if self.args.clear_quantized_weight_workspaces_on_offload:
+            num_cleared = clear_quantized_weight_workspaces(self.model)
+            logger.info(f"Dropped {num_cleared} cached quantized weight workspaces before offload")
+
         clear_memory(clear_host_memory=True)
         print_memory("before offload model")
         should_log_cpu_memory = is_first_replica_megatron_main_rank() and hasattr(self, "_last_rollout_id")

miles/utils/arguments.py

@@ -159,6 +159,16 @@ def add_cluster_arguments(parser):
                     "Defaults to true when --colocate is set; an explicit --no-offload-train is respected."
                 ),
             )
+            parser.add_argument(
+                "--clear-quantized-weight-workspaces-on-offload",
+                action=argparse.BooleanOptionalAction,
+                default=True,
+                help=(
+                    "Drop TransformerEngine's cached quantized weights before offloading the "
+                    "training actor. They are rebuilt on the next forward, so backing them up "
+                    "to pinned host memory is pure overhead. Unsupported with CUDA graphs."
+                ),
+            )
             parser.add_argument(
                 "--offload-rollout",
                 action=argparse.BooleanOptionalAction,

miles/utils/memory_utils.py

@@ -1,12 +1,37 @@
 import gc
 import logging
+from collections.abc import Sequence
 
 import torch
 import torch.distributed as dist
 
 logger = logging.getLogger(__name__)
 
 
+def clear_quantized_weight_workspaces(models: Sequence[torch.nn.Module]) -> int:
+    """Drop TransformerEngine's cached quantized weights so offload does not back them up.
+
+    Under a low-precision recipe TE keeps the quantized weight of every module in
+    ``_fp8_workspaces`` for both mxfp8, nvfp4. It is derived from the high-precision weight and TE rebuilds
+    it on the next forward via the cache-miss path.
+    """
+    from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
+
+    for model_chunk in models:
+        assert model_chunk.config.cuda_graph_impl == "none", (
+            "clear_quantized_weight_workspaces is unsafe with CUDA graphs "
+            f"(cuda_graph_impl={model_chunk.config.cuda_graph_impl})"
+        )
+
+    num_cleared = 0
+    for model_chunk in models:
+        for module in model_chunk.modules():
+            if isinstance(module, TransformerEngineBaseModule):
+                num_cleared += len(module._fp8_workspaces)
+                module._fp8_workspaces.clear()
+    return num_cleared
+
+
 def clear_memory(clear_host_memory: bool = False):
     torch.cuda.synchronize()
     gc.collect()