Fix NCCL out-of-memory error even when there is memory (THUDM#669)

Author: fzyzcjy

slime/ray/train_actor.py

@@ -6,6 +6,7 @@
 import ray
 import torch
 import torch.distributed as dist
+from torch_memory_saver import torch_memory_saver
 
 import slime.utils.eval_config
 from slime.ray.ray_actor import RayActor
@@ -46,6 +47,11 @@ def init(self, args, role, wandb_run_id, with_ref=False):
         self.role = role
         self.with_ref = with_ref
 
+        if (x := args.train_memory_margin_bytes) > 0:
+            print(f"Set torch_memory_saver.memory_margin_bytes to {x}")
+            assert args.offload_train
+            torch_memory_saver.memory_margin_bytes = x
+
         torch.serialization.add_safe_globals([slime.utils.eval_config.EvalDatasetConfig])
 
         local_rank = int(os.environ.get("LOCAL_RANK", 0))

slime/utils/arguments.py

@@ -130,6 +130,12 @@ def add_train_arguments(parser):
                 default="{}",
                 help="Extra environment variables for training process, e.g. PyTorch memory management ones.",
             )
+            parser.add_argument(
+                "--train-memory-margin-bytes",
+                type=int,
+                default=0,
+                help="Add margin for train memory allocation.",
+            )
 
             return parser