Support zero host or device memory waste for weight update (#276)

miles/backends/megatron_utils/actor.py

@@ -91,7 +91,10 @@ class MegatronTrainRayActor(TrainRayActor):
 
         self.weights_backuper = TensorBackuper.create(
             source_getter=lambda: named_params_and_buffers(
-                self.args, self.model, convert_to_global_name=args.megatron_to_hf_mode == "raw"
+                self.args,
+                self.model,
+                convert_to_global_name=args.megatron_to_hf_mode == "raw",
+                translate_gpu_to_cpu=not self.args.enable_weights_backuper,
             ),
             single_tag=None if args.enable_weights_backuper else "actor",
         )

miles/backends/megatron_utils/update_weight/common.py

@@ -117,11 +117,26 @@ def named_params_and_buffers(
     args: Namespace,
     model: Sequence[torch.nn.Module],
     convert_to_global_name: bool = True,
+    translate_gpu_to_cpu: bool = False,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     if convert_to_global_name:
-        return _named_params_and_buffers_global(args, model)
+        ans = _named_params_and_buffers_global(args, model)
     else:
-        return _named_params_and_buffers_vanilla(model)
+        ans = _named_params_and_buffers_vanilla(model)
+
+    if translate_gpu_to_cpu:
+        ans = ((name, _maybe_get_cpu_backup(tensor)) for name, tensor in ans)
+
+    return ans
+
+
+def _maybe_get_cpu_backup(x: torch.Tensor):
+    from torch_memory_saver import torch_memory_saver
+
+    if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None:
+        return cpu_tensor
+
+    return x
 
 
 def _named_params_and_buffers_vanilla(model: Sequence[torch.nn.Module]) -> Iterator[tuple[str, torch.Tensor]]:

train.py

@@ -28,12 +28,8 @@ def train(args):
     if args.offload_rollout:
         ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
 
-    if args.offload_train and not args.enable_weights_backuper:
-        actor_model.onload()
     # always update weight first so that sglang has the loaded weights from training.
     actor_model.update_weights()
-    if args.offload_train and not args.enable_weights_backuper:
-        actor_model.offload()
 
     if args.check_weight_update_equal:
         ray.get(rollout_manager.check_weights.remote(action="compare"))
@@ -93,15 +89,9 @@ def train(args):
             if args.rollout_global_dataset:
                 ray.get(rollout_manager.save.remote(rollout_id))
 
-        if args.enable_weights_backuper:
-            offload_train()
-            onload_rollout()
-            actor_model.update_weights()
-        else:
-            actor_model.clear_memory()
-            onload_rollout()
-            actor_model.update_weights()
-            offload_train()
+        offload_train()
+        onload_rollout()
+        actor_model.update_weights()
 
         if args.offload_rollout:
             if GPU_MEMORY_TYPE_CUDA_GRAPH is not None: