update code

Author: rumuru-bot

docker/Dockerfile

@@ -53,7 +53,7 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \
     cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \
     pip install -e .
 
-RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@9b8b788fdeb9c2ee528183214cef65a99b71e7d5 --no-cache-dir --force-reinstall
+RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall
 RUN pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
 RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation
 

scripts/run_glm45_355b_a32b.py

@@ -25,9 +25,13 @@ class ScriptArgs(U.ExecuteTrainConfig):
     enable_eval: bool = True
     extra_args: str = ""
     rollout_fp8: bool = False
+    rollout_attn_fp8: bool = False
     enable_mtp: bool = False  # TODO enable by default
     dynamic_sampling: bool = False
     enable_benchmark: bool = False
+    enable_mis: bool = False
+    # TODO improve, should be able to override more easily
+    tis_use_rs: bool = True
     task: Literal["dapo_aime", "gsm8k"] = "dapo_aime"
 
 
@@ -243,9 +247,11 @@ def train(args: ScriptArgs):
         # """--sglang-json-model-override-args '{"num_hidden_layers": 5}' """
     )
     sglang_extra_env_vars = {}
+    if U.GENERATION_HARDWARE[args.hardware] == "Blackwell":
+        sglang_args += "--sglang-attention-backend trtllm_mha "
     if args.rollout_fp8:
         sglang_decode_max_bs = 256
-        sglang_attn_tp_size = 8
+        sglang_attn_tp_size = min(8, sglang_world_size)
         sglang_attn_dp_size = sglang_world_size // sglang_attn_tp_size
         sglang_args += (
             f"--sglang-ep-size {sglang_world_size} "
@@ -306,6 +312,35 @@ def train(args: ScriptArgs):
     if args.enable_benchmark:
         misc_args += (
             "--custom-generate-function-path slime.rollout.generate_hub.benchmarkers.generate_with_random_osl "
+            "--rollout-batch-size 128 "
+            "--n-samples-per-prompt 8 "
+            "--use-distributed-post "
+            "--router-policy round_robin "
+            "--sglang-server-concurrency 10000 "
+            # GB200 w/ mem-frac 0.8 will lead to oom in long jobs currently, but here we use large value to make baseline more fair
+            f"--sglang-mem-fraction-static {0.8 if args.hardware == 'GB300' else 0.75} "
+        )
+
+    if args.rollout_attn_fp8:
+        sglang_args += "--sglang-kv-cache-dtype fp8_e4m3 "
+
+    if args.enable_mis:
+        config_text = f"""
+use_tis: true
+use_rs: {"true" if args.tis_use_rs else "false"}
+tis_level: "token"
+rs_level: "token"
+tis_mode: "truncate"
+tis_lower_bound: 0.5
+tis_upper_bound: 2.0
+rs_lower_bound: null
+rs_upper_bound: null
+rs_veto_threshold: 1.0e-4
+tis_batch_normalize: true
+""".strip()
+        misc_args += (
+            f"--custom-config-path {U.save_to_temp_file(config_text, 'yaml')} "
+            "--custom-tis-function-path examples.train_infer_mismatch_helper.mis.compute_mis_weights_with_cp "
         )
 
     train_args = (

slime/backends/sglang_utils/sglang_engine.py

@@ -6,7 +6,6 @@
 import requests
 import sglang_router
 from packaging.version import parse
-from sglang.srt.entrypoints.http_server import launch_server
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import kill_process_tree
 from urllib3.exceptions import NewConnectionError
@@ -31,6 +30,8 @@ def get_base_gpu_id(args, rank):
 
 
 def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
+    from sglang.srt.entrypoints.http_server import launch_server
+
     multiprocessing.set_start_method("spawn", force=True)
     server_args.host = server_args.host.strip("[]")
     p = multiprocessing.Process(target=launch_server, args=(server_args,))

slime/rollout/rm_hub/__init__.py

@@ -1,4 +1,5 @@
 import asyncio
+import random
 
 import aiohttp
 
@@ -57,6 +58,8 @@ async def async_rm(args, sample: Sample, **kwargs):
         from .ifbench import compute_ifbench_reward
 
         return compute_ifbench_reward(response, label, metadata=metadata)
+    elif rm_type == "random":
+        return random.randint(0, 1)
     elif rm_type:
         raise NotImplementedError(f"Rule-based RM for {rm_type} is not implemented.")
     else:

slime/rollout/sglang_rollout.py

@@ -293,9 +293,8 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]:
         response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
         urls = [worker["url"] for worker in response["workers"]]
 
-    for url in urls:
-        logger.info(f"Abort request for {url}")
-        await post(f"{url}/abort_request", {"abort_all": True})
+    logger.info(f"Abort request for {urls}")
+    await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls])
 
     # make sure all the pending tasks are finished
     count = 0

slime/utils/external_utils/command_utils.py

@@ -262,3 +262,9 @@ def save_to_temp_file(text: str, ext: str):
     "GB200": 4,
     "GB300": 4,
 }
+
+GENERATION_HARDWARE = {
+    "H100": "Hopper",
+    "GB200": "Blackwell",
+    "GB300": "Blackwell",
+}

slime/utils/tensor_backper.py

@@ -31,7 +31,6 @@ def get(self, tag: str):
     def backup(self, tag: str):
         raise NotImplementedError
 
-    @abstractmethod
     def copy(self, *, src_tag: str, dst_tag: str):
         raise NotImplementedError