Fix residual GPU memory usage in TrainingWorker in RL workflow

examples/v1/config/reasoning_rl_qwen3p5vl_mtp_ep.py

@@ -30,6 +30,7 @@
 import ray
 
 # export LMDEPLOY_FP32_MAMBA_SSM_DTYPE=1
+# export XTUNER_ACTIVATION_OFFLOAD=1
 
 def _as_list(value):
     return value if isinstance(value, list) else [value]
@@ -61,6 +62,7 @@ def _as_list(value):
 checkpoint_interval = 50
 evaluate_step = 5
 enable_initial_evaluate = os.environ.get("ENABLE_INITIAL_EVALUATE", False)
+train_ep_size = 4
 
 # 1. resources
 resources = AcceleratorResourcesConfig(
@@ -89,7 +91,7 @@ def _as_list(value):
         lmdeploy_uvicorn_log_level="INFO",
         lmdeploy_speculative_algorithm='qwen3_5_mtp',
         # MTP draft tokens trade throughput for extra activation memory; try 3 if still tight.
-        lmdeploy_speculative_num_draft_tokens=4,
+        lmdeploy_speculative_num_draft_tokens=3,
     ),
     health_check_interval_seconds=300,
     health_check_failure_threshold=3,
@@ -207,15 +209,21 @@ def _as_list(value):
     },
 )
 
+from xtuner.v1.float8 import Float8Config, ScalingGranularity
 # 5. train worker
 model_cfg = Qwen3_5_VLMoE35BA3Config(freeze_vision=True, freeze_projector=True)
-model_cfg.float8_cfg = None
-model_cfg.text_config.ep_size = 1
+float8_cfg = Float8Config(
+    scaling_granularity_gemm=None,
+    scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE,
+)
+model_cfg.float8_cfg = float8_cfg
+model_cfg.text_config.ep_size = train_ep_size
 model_cfg.text_config.z_loss_cfg = None
 model_cfg.text_config.balancing_loss_cfg = None
 model_cfg.text_config.freeze_routers = True
+model_cfg.compile_cfg = None
 model_cfg.text_config.mtp_config = MTPConfig(
-    num_layers=4, 
+    num_layers=3, 
     loss_scaling_factor=1.0,
     detach_mtp_lm_head_weight=True,
     detach_mtp_inputs=True,
@@ -246,7 +254,7 @@ def _as_list(value):
     ),
 )
 lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6)
-fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=1, fp32_lm_head=True)
+fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=train_ep_size, fp32_lm_head=True)
 train_worker_cfg = WorkerConfig(
     model_cfg=model_cfg,
     load_from=model_path,

setup.cfg

@@ -1,5 +1,5 @@
 [codespell]
-ignore-words-list = nd, ba, warmup, ans, dout
+ignore-words-list = nd, ba, warmup, ans, dout, ptd
 
 [flake8]
 max-line-length = 119

xtuner/v1/rl/trainer/controller.py

@@ -1,6 +1,6 @@
 import math
 import os
-from typing import Literal, TypedDict
+from typing import Any, Literal, TypedDict
 
 import ray
 import torch
@@ -24,6 +24,23 @@ class ColateItem(TypedDict):
     rollout_logprobs: torch.Tensor | None
 
 
+def _summarize_process_group_results(results: list[dict[str, Any]]) -> str:
+    if not results:
+        return "ranks=0"
+
+    count_key = "destroyed" if "destroyed" in results[0] else "reloaded"
+    counts = [result.get(count_key, 0) for result in results]
+    count_summary = f"{counts[0]} on all ranks" if len(set(counts)) == 1 else f"by_rank={counts}"
+    result_errors = [error for result in results for error in result.get("errors", [])]
+    unwrapped = sum(len(result.get("unwrapped_nccl_groups", [])) for result in results)
+    summary = f"ranks={len(results)}, {count_key}={count_summary}"
+    if unwrapped:
+        summary += f", unwrapped_nccl_groups={unwrapped}"
+    if result_errors:
+        summary += f", errors={len(result_errors)}, first_error={result_errors[0]}"
+    return summary
+
+
 class TrainingController:
     def __init__(self, workers: list[TrainingWorker]) -> None:
         self.workers = workers
@@ -246,8 +263,6 @@ def fit(self, data_batches: list[ColateItem], pack_max_length: int, rollout_idx:
             pad_data_samples = [pad_data for _ in range(pad_num)]
             packed_data_batches = packed_data_batches + pad_data_samples
 
-        print(f"len(packed_data_batches): {len(packed_data_batches)}")
-
         handles = []
         data_batch_refs = {}
         for worker_idx, worker in enumerate(self.workers):
@@ -323,6 +338,26 @@ def update_weights(self):
         ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)
         return
 
+    def destroy_train_nccl_process_groups(self):
+        """Destroy train-side NCCL process groups after weight sync."""
+        handles = [
+            worker.destroy_train_nccl_process_groups.remote()  # type: ignore[attr-defined]
+            for worker in self.workers
+        ]
+        results = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)
+        self.logger.info(f"Destroyed train NCCL process groups: {_summarize_process_group_results(results)}")
+        return results
+
+    def reload_train_nccl_process_groups(self):
+        """Reload train-side reloadable NCCL process groups before training."""
+        handles = [
+            worker.reload_train_nccl_process_groups.remote()  # type: ignore[attr-defined]
+            for worker in self.workers
+        ]
+        results = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)
+        self.logger.info(f"Reloaded train NCCL process groups: {_summarize_process_group_results(results)}")
+        return results
+
     def save_hf(self, hf_dir: str, save_dtype: torch.dtype = torch.bfloat16):
         handles = [worker.save_hf.remote(hf_dir, save_dtype) for worker in self.workers]  # type: ignore
         ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)

xtuner/v1/rl/trainer/worker.py

@@ -57,6 +57,14 @@
     ray_method,
     set_deterministic,
 )
+from xtuner.v1.utils.activation_offload import OffloadManager
+from xtuner.v1.utils.fsdp import release_deferred_fsdp_all_gathers
+from xtuner.v1.utils.reloadable_process_group import (
+    destroy_reloadable_process_groups,
+    monkey_patch_reloadable_process_groups,
+    process_group_name,
+    reload_process_groups,
+)
 
 from ..rollout_is import merge_rollout_is_metrics
 
@@ -221,6 +229,7 @@ def __init__(
         self.config = cast(WorkerConfig, self.config)
         torch.accelerator.set_device_index(int(os.environ["LOCAL_RANK"]))
         self.rank = rank
+        monkey_patch_reloadable_process_groups()
 
         # TODO: add lr scheduler
         log_dir = worker_cfg.log_dir
@@ -278,6 +287,74 @@ def __init__(
             engine=self._engine,
         )
 
+    @ray_method
+    def destroy_train_nccl_process_groups(self) -> dict[str, object]:
+        """Destroy reloadable train-side NCCL PG inners after weight sync and
+        train offload."""
+
+        details = destroy_reloadable_process_groups()
+        errors: list[str] = []
+        unwrapped_nccl_groups: list[dict[str, object]] = []
+
+        try:
+            import torch.distributed.distributed_c10d as c10d
+        except Exception as exc:
+            errors.append(f"import distributed_c10d {type(exc).__name__}: {exc}")
+            c10d = None  # type: ignore
+
+        groups: list[tuple[dist.ProcessGroup, list[int], str]] = []
+        default_group = dist.group.WORLD if dist.is_initialized() else None
+        if c10d is not None:
+            for group, rank_map in list(c10d._world.pg_group_ranks.items()):
+                if group is default_group:
+                    continue
+                try:
+                    backend = str(dist.get_backend(group)).lower()
+                except ValueError:
+                    # The c10d registry may still contain entries for groups
+                    # just destroyed by the reloadable wrapper. They are not
+                    # actionable unwrapped NCCL groups.
+                    continue
+                except Exception as exc:
+                    errors.append(f"get_backend {type(exc).__name__}: {exc}")
+                    continue
+                if "nccl" not in backend:
+                    continue
+                ranks = [global_rank for global_rank, _ in sorted(rank_map.items(), key=lambda item: item[1])]
+                groups.append((group, ranks, backend))
+
+        for group, ranks, backend in groups:
+            if group.__class__.__name__ == "ReloadableProcessGroup":
+                continue
+            unwrapped_nccl_groups.append({"ranks": ranks, "backend": backend, "group_name": process_group_name(group)})
+
+        DEVICE_MODULE.empty_cache()
+        result = {
+            "rank": self.rank,
+            "destroyed": len(details),
+            "unwrapped_nccl_groups": unwrapped_nccl_groups,
+            "errors": errors,
+        }
+        return result
+
+    @ray_method
+    def reload_train_nccl_process_groups(self) -> dict[str, object]:
+        """Reload train-side NCCL PG inners before FSDP2 runs again."""
+
+        errors: list[str] = []
+        try:
+            details = reload_process_groups()
+        except Exception as exc:
+            details = []
+            errors.append(f"reload {type(exc).__name__}: {exc}")
+        DEVICE_MODULE.empty_cache()
+        result = {
+            "rank": self.rank,
+            "reloaded": len(details),
+            "errors": errors,
+        }
+        return result
+
     @ray_method
     def bind_rollout_weight_update(self, *args, **kwargs):
         return self.update_weighter.bind_rollout_weight_update(*args, **kwargs)
@@ -679,6 +756,13 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo
 
         only_calc_mismatch_ratio = os.environ.get("ONLY_CALC_MISMATCH_RATIO", "0") == "1"
         if only_calc_mismatch_ratio:
+            # This early-return path skips the MTP forward. FSDP may have
+            # prefetched the MTP module already, leaving its all-gather result
+            # cached until we release it explicitly. The normal training branch
+            # can leave the same residue, but it is outside the peak-memory
+            # path and can be ignored.
+            self._release_deferred_fsdp_all_gathers(f"rollout_{rollout_idx}/only_calc_mismatch")
+            DEVICE_MODULE.empty_cache()
             return worker_log_item
 
         # compute reference logprobs
@@ -772,6 +856,7 @@ def fit(self, data_batches: list[WorkerInputItem], rollout_idx: int) -> WorkerLo
                 train_step_info = self._engine.train_step(
                     data_batches=engine_input,
                 )
+                OffloadManager().clear()
             self.logger.debug(
                 f"Rank{self.rank} Rollout {rollout_idx} GlobalStep {global_train_step} "
                 f"train_step[{i}].engine_train_step elapsed={time.perf_counter() - train_step_begin:.4f}s"
@@ -954,10 +1039,30 @@ def get_model_cfg(self):
         model_cfg = self._engine.model_cfg
         return model_cfg
 
+    def _clear_cublas_workspaces(self) -> None:
+        clear_fn = getattr(torch._C, "_cuda_clearCublasWorkspaces", None)
+        if clear_fn is None:
+            self.logger.warning("torch._C._cuda_clearCublasWorkspaces is unavailable")
+            return
+        try:
+            clear_fn()
+            DEVICE_MODULE.empty_cache()
+        except Exception as e:
+            self.logger.warning(f"Failed to clear cuBLAS workspaces: {e}")
+
+    def _release_deferred_fsdp_all_gathers(self, log_tag: str) -> None:
+        comm_states, param_groups = release_deferred_fsdp_all_gathers(self._engine.model)
+        if comm_states or param_groups:
+            self.logger.debug(
+                f"[{log_tag}] released deferred FSDP all-gathers: "
+                f"comm_states={comm_states}, param_groups={param_groups}"
+            )
+
     @ray_method
     def offload_model(self):
         self._engine.put_model_to_device("cpu")
         DEVICE_MODULE.empty_cache()
+        self._clear_cublas_workspaces()
         self.logger.info(
             f"Offloaded model to CPU. Current allocate {DEVICE_MODULE.memory_allocated() / (1024**2)} MB, reserved: {DEVICE_MODULE.memory_reserved() / (1024**2)} MB"
         )

xtuner/v1/train/rl_trainer.py

@@ -913,6 +913,12 @@ def _train_one_batch(
             )
             ray.get(self.rollout_controller.offload.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT)
         if onload_train_before_train:
+            if getattr(self, "_train_nccl_destroyed", False):
+                # FSDP2 keeps DeviceMesh/PG references alive. Recreate the
+                # wrapped NCCL inners before loading train weights again.
+                with timer("reload_train_nccl", step_timer_dict):
+                    self.train_controller.reload_train_nccl_process_groups()
+                self._train_nccl_destroyed = False
             with timer("onload", step_timer_dict):
                 self.train_controller.onload(target="all")
                 self.logger.info("Training controller loaded")
@@ -1772,6 +1778,15 @@ def _sync_weights_and_save(self, train_step: int, step_timer_dict: dict) -> bool
                 self.train_controller.update_weights()
                 self.logger.info("Rollout workers update weights successfully in colocate mode")
                 self.train_controller.offload(target="model")
+                destroy_train_nccl = os.getenv("XTUNER_DESTROY_TRAIN_NCCL_AFTER_SYNC", "0") == "1"
+                already_destroyed = getattr(self, "_train_nccl_destroyed", False)
+                if destroy_train_nccl and not already_destroyed:
+                    # This runs only after rollout weights have been updated
+                    # and the train model has been offloaded. Only replayable
+                    # non-default NCCL groups are destroyed.
+                    with timer("destroy_train_nccl", step_timer_dict):
+                        self.train_controller.destroy_train_nccl_process_groups()
+                    self._train_nccl_destroyed = True
             else:
                 self.train_controller.offload(target="model")
                 ray.get(self.rollout_controller.onload_weights.remote(), timeout=RL_TRAINER_RAY_GET_TIMEOUT)

xtuner/v1/utils/activation_offload.py

@@ -287,12 +287,42 @@ def prefetch_get(self, block_idx, tensor_idx, h2d_stream, d2h_stream, group="def
     def empty(self):
         return len(self.items) == 0
 
-    def clear(self, key=None):
-        if key is None:
-            self.items.clear()
-        else:
-            self.assert_exist(key)
-            self.items.pop(key)
+    def clear(self, key=None, group=None, clear_pin_memory_cache: bool = False):
+        """Clear runtime offload state.
+
+        ``pin_memory_cache`` is kept by default because it is a reusable CPU
+        pinned-memory cache, not a live GPU activation reference.
+        """
+        if key is not None and group is not None:
+            raise ValueError("Only one of key or group can be specified")
+
+        if key is not None:
+            self.items.pop(key, None)
+            self.may_npu_tensors.pop(key, None)
+            if clear_pin_memory_cache:
+                self.pin_memory_cache.pop(key, None)
+            return
+
+        if group is not None:
+            prefix = f"{group}_"
+            for item_key in list(self.items.keys()):
+                if item_key.startswith(prefix):
+                    self.items.pop(item_key, None)
+            for item_key in list(self.may_npu_tensors.keys()):
+                if item_key.startswith(prefix):
+                    self.may_npu_tensors.pop(item_key, None)
+            self.getcnt.pop(group, None)
+            if clear_pin_memory_cache:
+                for item_key in list(self.pin_memory_cache.keys()):
+                    if item_key.startswith(prefix):
+                        self.pin_memory_cache.pop(item_key, None)
+            return
+
+        self.items.clear()
+        self.may_npu_tensors.clear()
+        self.getcnt.clear()
+        if clear_pin_memory_cache:
+            self.pin_memory_cache.clear()
 
     # event interface #