sync from internal

slime/backends/megatron_utils/arguments.py

@@ -1,5 +1,4 @@
 import logging
-
 from megatron.training.arguments import parse_args as _megatron_parse_args
 from megatron.training.arguments import validate_args as _megatron_validate_args
 from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding

slime/backends/megatron_utils/model.py

@@ -725,6 +725,7 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None:
 
     try:
         from megatron.bridge import AutoBridge
+
         from slime.utils.megatron_bridge_utils import patch_megatron_model
 
         path = Path(args.save_hf.format(rollout_id=rollout_id))
@@ -765,6 +766,7 @@ def initialize_model_and_optimizer(
 
     if torch.version.hip:
         import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
+
         from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
 
         filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync

slime/backends/megatron_utils/model_provider.py

@@ -140,7 +140,17 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage
             transformer_layer_spec = import_module(args.spec)
             # Allow the spec to be a function so that user can use customized Megatron easier.
             if callable(transformer_layer_spec):
-                transformer_layer_spec = transformer_layer_spec(args, config, vp_stage)
+                result = transformer_layer_spec(args, config, vp_stage)
+                # If the result is itself a model provider (callable with pre_process param),
+                # delegate model construction to it directly (e.g. glm-omni VL model).
+                if callable(result) and "pre_process" in inspect.signature(result).parameters:
+                    model = result(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
+                    if post_process and role == "critic":
+                        model.output_layer = LinearForLastLayer(
+                            input_size=config.hidden_size, output_size=1, config=config
+                        )
+                    return model
+                transformer_layer_spec = result
         else:
             if args.num_experts:
                 # Define the decoder block spec

slime/backends/megatron_utils/update_weight/common.py

@@ -37,7 +37,7 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
     assert param.partition_stride == 1, "partition_stride != 1 is not supported"
     # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
     # TODO: check only GLU is used.
-    if "linear_fc1.weight" in name:
+    if "linear_fc1.weight" in name or "linear_fc1.bias" in name:
         param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
         param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
     # this is bug in megatron's grouped moe.
@@ -99,7 +99,7 @@ def all_gather_params_async(
             assert partition_dim is not None, "partition_stride != 1 is not supported"
             # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
             # TODO: check only GLU is used.
-            if "linear_fc1.weight" in info.name:
+            if "linear_fc1.weight" in info.name or "linear_fc1.bias" in info.name:
                 param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
                 param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
             # this is bug in megatron's grouped moe.

slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py

@@ -170,7 +170,6 @@ def _get_megatron_local_param_infos(args: Namespace, model: Sequence[torch.nn.Mo
                 continue
             for name, info in infos.items():
                 if name in param_infos:
-                    assert args.mtp_num_layers is not None
                     old_info = param_infos[name]
                     if old_info.src_rank > src_rank:
                         param_infos[name] = info

slime/backends/sglang_utils/sglang_engine.py

@@ -52,17 +52,23 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int:
 
 def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
     if getattr(server_args, "encoder_only", False):
-        from sglang.srt.disaggregation.encode_server import launch_server
-    else:
-        from sglang.srt.entrypoints.http_server import launch_server
+        from sglang.srt.disaggregation.encode_server import launch_server_process as sglang_launch_server_process
+
+        return sglang_launch_server_process(
+            server_args,
+            start_method="spawn",
+            wait_for_server=True,
+        )
+
+    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,))
     p.start()
 
-    if server_args.node_rank != 0:
-        return
+    if getattr(server_args, "node_rank", 0) != 0:
+        return p
 
     _wait_server_healthy(
         base_url=server_args.url(),
@@ -198,7 +204,7 @@ def _init_normal(self, server_args_dict):
             if parse(sglang_router.__version__) <= parse("0.2.1"):
                 assert self.worker_type == "regular", "pd disaggregation is not supported in old router."
                 response = requests.post(
-                    f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}"
+                    f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}",
                 )
             else:
                 payload = {

slime/ray/rollout.py

@@ -719,9 +719,13 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
             loss_masks.append(sample.loss_mask)
         train_data["loss_masks"] = loss_masks
 
-        # overwriting the raw reward
-        if samples[0].metadata and "raw_reward" in samples[0].metadata:
-            train_data["raw_reward"] = [sample.metadata["raw_reward"] for sample in samples]
+        # Overwrite raw_reward when available. Mixed-source batches may only
+        # populate this field for a subset of samples (e.g. SWE but not code).
+        if any(sample.metadata and "raw_reward" in sample.metadata for sample in samples):
+            train_data["raw_reward"] = [
+                sample.metadata["raw_reward"] if sample.metadata and "raw_reward" in sample.metadata else sample.reward
+                for sample in samples
+            ]
 
         # For rollout buffer
         if samples[0].metadata and "round_number" in samples[0].metadata:
@@ -1076,13 +1080,15 @@ def _make_group(group_cfg, router_ip, router_port, overrides_extra=None):
 
             logger.info(f"EPD phase 1 done: collected {len(encoder_urls)} encoder URLs: {encoder_urls}")
 
-            # --- Phase 2: start non-encoder groups, injecting encoder URLs into prefill ---
+            # --- Phase 2: start non-encoder groups, injecting encoder URLs into
+            # language-only LLM workers. Prefill groups use this for full EPD,
+            # while regular groups allow encoder/LLM split without PD.
             non_encoder_handles: list = []
             for group_cfg in model_cfg.server_groups:
                 if group_cfg.worker_type == "encoder":
                     continue
                 overrides_extra = {}
-                if encoder_urls and group_cfg.worker_type == "prefill":
+                if encoder_urls and group_cfg.worker_type in ("prefill", "regular"):
                     overrides_extra["language_only"] = True
                     overrides_extra["encoder_urls"] = encoder_urls
                 group = _make_group(group_cfg, router_ip, router_port, overrides_extra=overrides_extra)

slime/rollout/sglang_rollout.py

@@ -5,6 +5,7 @@
 import uuid
 from argparse import Namespace
 from collections.abc import Callable
+from contextlib import contextmanager
 from typing import Any
 
 import numpy as np
@@ -35,6 +36,30 @@
 
 logger = logging.getLogger(__name__)
 
+_PROCESSOR_PROMPT_KEYS = {"input_ids", "attention_mask"}
+
+
+def _prepare_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]:
+    raw_multimodal_inputs = sample.multimodal_inputs or {}
+    has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values())
+    reuse_existing_input_ids = bool(sample.tokens) and (
+        sample.multimodal_train_inputs is not None or not has_multimodal_inputs
+    )
+
+    if processor and has_multimodal_inputs and not reuse_existing_input_ids:
+        processor_output = processor(text=sample.prompt, **build_processor_kwargs(raw_multimodal_inputs))
+        prompt_ids = processor_output["input_ids"][0]
+        if sample.multimodal_train_inputs is None:
+            sample.multimodal_train_inputs = {
+                k: v for k, v in processor_output.items() if k not in _PROCESSOR_PROMPT_KEYS
+            } or None
+        return prompt_ids
+
+    if reuse_existing_input_ids:
+        return sample.tokens
+
+    return tokenizer.encode(sample.prompt, add_special_tokens=False)
+
 
 def get_model_url(args: Namespace, model_name: str, endpoint: str = "/generate") -> str:
     """Return the router URL for a named model.
@@ -85,8 +110,24 @@ def __init__(self, args: Namespace) -> None:
             sampling_seed_base = args.rollout_seed
             self.group_sampling_seeds = [sampling_seed_base + i for i in range(args.n_samples_per_prompt)]
 
+        # dp rank balancing
+        self.dp_counts = [0] * (args.sglang_dp_size or 1)
+        self.dp_rank = 0
+
         self.reset()
 
+    @contextmanager
+    def dp_rank_context(self):
+        candidates = [i for i, count in enumerate(self.dp_counts) if count == min(self.dp_counts)]
+        dp_rank = int(np.random.choice(candidates))
+        self.dp_counts[dp_rank] += 1
+        self.dp_rank = dp_rank
+        try:
+            yield dp_rank
+        finally:
+            self.dp_counts[dp_rank] -= 1
+            assert self.dp_counts[dp_rank] >= 0
+
     def reset(self) -> None:
         self.remaining_batch_size = 0
         self.pendings = set()
@@ -120,18 +161,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED
     ), f"Sample status is {sample.status}"
 
-    if state.processor and sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()):
-        processor_kwargs = build_processor_kwargs(sample.multimodal_inputs)
-        processor_output = state.processor(text=sample.prompt, **processor_kwargs)
-        prompt_ids = processor_output["input_ids"][0]
-        sample.multimodal_train_inputs = {
-            k: v for k, v in processor_output.items() if k not in ["input_ids", "attention_mask"]
-        } or None
-    else:
-        prompt_ids = state.tokenizer.encode(sample.prompt, add_special_tokens=False)
-
-    if len(sample.response) > 0:
-        sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids)
+    prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor)
 
     assert (
         sampling_params["max_new_tokens"] >= 0
@@ -149,25 +179,17 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
     if args.use_rollout_routing_replay:
         payload["return_routed_experts"] = True
 
-    has_multimodal = sample.multimodal_inputs and sample.multimodal_inputs.get("images")
-    if has_multimodal:
-        image_data = sample.multimodal_inputs["images"]
-        payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data]
-
-    # Use existing tokens for multi-turn or tokenize the new prompt
-    if len(sample.response) > 0:
-        payload["input_ids"] = sample.tokens
-    elif has_multimodal:
-        # For multimodal first-turn: send text so SGLang handles image token
-        # expansion internally (the processor-expanded input_ids have N patch
-        # tokens per image which would mismatch the image_data count).
+    images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None
+    if images:
+        payload["image_data"] = [encode_image_for_rollout_engine(image) for image in images]
+        # For single-turn multimodal requests, send text so SGLang expands the
+        # image placeholders with its own processor rules.
         payload["text"] = sample.prompt
-        if not sample.tokens:
-            sample.tokens = prompt_ids
     else:
         payload["input_ids"] = prompt_ids
-        if not sample.tokens:  # Initialize sample.tokens for the first turn
-            sample.tokens = prompt_ids
+
+    if not sample.tokens:
+        sample.tokens = prompt_ids
 
     # Use session_id for consistent hashing routing (SGLang Model Gateway)
     headers = None
@@ -240,30 +262,29 @@ async def generate_and_rm(
             sample.status = Sample.Status.ABORTED
             return sample
 
-        # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config)
-        custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path
-
-        if custom_func_path is not None:
-            custom_generate_func = load_function(custom_func_path)
-            # if signature has evaluation, pass evaluation
-            if "evaluation" in inspect.signature(custom_generate_func).parameters:
-                sample = await custom_generate_func(args, sample, sampling_params, evaluation=evaluation)
+        with state.dp_rank_context() as _:
+            # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config)
+            custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path
+
+            if custom_func_path is not None:
+                custom_generate_func = load_function(custom_func_path)
+                # if signature has evaluation, pass evaluation
+                if "evaluation" in inspect.signature(custom_generate_func).parameters:
+                    sample = await custom_generate_func(args, sample, sampling_params, evaluation=evaluation)
+                else:
+                    sample = await custom_generate_func(args, sample, sampling_params)
             else:
-                sample = await custom_generate_func(args, sample, sampling_params)
-        else:
-            sample = await generate(args, sample, sampling_params)
+                sample = await generate(args, sample, sampling_params)
 
     # for the rm that need the whole group, we will not do the rm here
     if args.group_rm:
         return sample
 
-    # multi samples
     if isinstance(sample, list):
         samples = sample
-        if any([sample.status == Sample.Status.ABORTED for sample in samples]):
+        if any(sample.status == Sample.Status.ABORTED for sample in samples):
             return samples
 
-        # for multi agent system, the reward of some sample is calculated during generation.
         samples_need_reward = [sample for sample in samples if sample.reward is None]
         with trace_span(samples_need_reward, "reward_model"):
             rewards = await batched_async_rm(args, samples_need_reward)
@@ -273,7 +294,7 @@ async def generate_and_rm(
     else:
         if sample.status == Sample.Status.ABORTED:
             return sample
-        # for multi-turn environment, a reward could be assigned to the agent.
+        # Some custom generate paths may have already filled the reward.
         if sample.reward is None:
             with trace_span(sample, "reward_model"):
                 sample.reward = await async_rm(args, sample)
@@ -597,5 +618,6 @@ def generate_rollout(
         return output
 
     output, aborted_samples = run(generate_rollout_async(args, rollout_id, data_source.get_samples))
-    data_source.add_samples(aborted_samples)
+    if aborted_samples:
+        data_source.add_samples(aborted_samples)
     return output