From f717332dcb6d3b30a5343aeaba317407d1b318a8 Mon Sep 17 00:00:00 2001
From: zhihaow6 <zhihaow6@illinois.edu>
Date: Tue, 4 Aug 2026 00:13:30 -0700
Subject: [PATCH] Drop TE cached quantized weights before offloading the
 training actor

Under a low-precision recipe TE caches each module's quantized weight in
_fp8_workspaces and never clears it. In colocate RL every offload copies
those workspaces to pinned host memory, even though they are derived from
the high-precision weights and are rebuilt on the next forward.

Clearing them before offload cuts host traffic per offload by 45% on
DeepSeek-V4-Flash MXFP8 (52.3 GB over 4 ranks) and 4.6 GB on GLM-5.2 NVFP4.

Gated by --clear-quantized-weight-workspaces-on-offload, asserted off
under CUDA graphs.
---
 miles/backends/megatron_utils/actor.py |  6 +++++-
 miles/utils/arguments.py               | 10 ++++++++++
 miles/utils/memory_utils.py            | 25 +++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py
index 58fe75a0f41..69f7859ddaa 100644
--- a/miles/backends/megatron_utils/actor.py
+++ b/miles/backends/megatron_utils/actor.py
@@ -22,7 +22,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
@@ -289,6 +289,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")
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 8d4318b83ba..295ebd816eb 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -154,6 +154,16 @@ def add_cluster_arguments(parser):
                     "This will always be true when --colocate is set."
                 ),
             )
+            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,
diff --git a/miles/utils/memory_utils.py b/miles/utils/memory_utils.py
index c12f3cd0bcd..427cac81ef7 100644
--- a/miles/utils/memory_utils.py
+++ b/miles/utils/memory_utils.py
@@ -1,5 +1,6 @@
 import gc
 import logging
+from collections.abc import Sequence
 
 import torch
 import torch.distributed as dist
@@ -7,6 +8,30 @@
 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()