[Feat] Support offload cuda graph (#354)

Author: ryang-max

slime/ray/rollout.py

@@ -272,6 +272,7 @@ def init_rollout_engines(args, pg, all_rollout_engines):
                     "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false",
                     "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true",
                     "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true",
+                    "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true",
                 }
             },
         ).remote(args, rank=i)

train.py

@@ -1,6 +1,11 @@
 import ray
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
+try:
+    from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
+except ImportError:
+    GPU_MEMORY_TYPE_CUDA_GRAPH = None
+
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
 from slime.utils.wandb_utils import init_wandb_primary
@@ -25,6 +30,8 @@ def train(args):
     actor_model.update_weights()
 
     if args.offload:
+        if GPU_MEMORY_TYPE_CUDA_GRAPH is not None:
+            ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]))
         ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
 
     # train loop.
@@ -70,6 +77,8 @@ def train(args):
         actor_model.update_weights()
 
         if args.offload:
+            if GPU_MEMORY_TYPE_CUDA_GRAPH is not None:
+                ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH]))
             ray.get(rollout_manager.onload.remote(tags=[GPU_MEMORY_TYPE_KV_CACHE]))
 
         if args.eval_interval is not None and (