sync internal features (THUDM#1192)

Author: zhuzilin

slime/backends/megatron_utils/actor.py

@@ -25,6 +25,7 @@
 from slime.utils.timer import Timer, inverse_timer, timer
 from slime.utils.tracking_utils import init_tracking
 from slime.utils.types import RolloutBatch
+
 from ...utils.profile_utils import TrainProfiler
 from ...utils.tensor_backper import TensorBackuper
 from .checkpoint import load_checkpoint
@@ -66,7 +67,6 @@ def init(
             if i == dist.get_rank():
                 self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
                 self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
-
             dist.barrier(group=get_gloo_group())
 
         self.train_parallel_config = {

slime/backends/megatron_utils/data.py

@@ -376,6 +376,64 @@ def log_rollout_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatc
     if args.log_passrate:
         log_passrate(rollout_id, args, rollout_data)
 
+    if args.log_correct_samples:
+        if mpu.get_tensor_model_parallel_rank() == 0 and mpu.is_pipeline_last_stage():
+            cp_size = mpu.get_context_parallel_world_size()
+            log_dict = {}
+            response_lengths = rollout_data["response_lengths"]
+            loss_masks = rollout_data["loss_masks"]
+            total_lengths = rollout_data["total_lengths"]
+
+            def quantile(total_value, n_quantiles, data) -> dict:
+                import math
+
+                assert n_quantiles > 1, f"n_quantiles({n_quantiles}) must be greater than 1."
+
+                quantiles = [((i + 1) / n_quantiles) for i in range(n_quantiles)]
+                cut_points = [total_value * q for q in quantiles]
+                cut_points[-1] = total_value
+
+                count = [0] * n_quantiles
+                for d in data:
+                    for i, point in enumerate(cut_points):
+                        if d <= point:
+                            count[i] += 1
+                            break
+
+                total = sum(count) + 1e-9
+                percentile = [c / total for c in count]
+
+                percentile = {f"p{min(math.ceil(q*100),100)}": p for q, p in zip(quantiles, percentile, strict=True)}
+                return percentile
+
+            raw_rewards = rollout_data["raw_reward"]
+            # Additional metrics for correct cases are calculated separately below.
+            correct_response_lengths = []
+            correct_total_lengths = []
+            correct_loss_masks = []
+            correct_entropy = []
+            for i, raw_reward in enumerate(raw_rewards):
+                if raw_reward == 1:
+                    correct_response_lengths.append(response_lengths[i])
+                    correct_total_lengths.append(total_lengths[i])
+                    correct_loss_masks.append(loss_masks[i])
+                    correct_entropy.append(-rollout_data["log_probs"][i])
+            num_correct_responses = len(correct_total_lengths)
+            rollout_data["correct_response_lengths"] = correct_response_lengths
+            correct_response_length_percentile = quantile(
+                args.rollout_max_response_len, 4, rollout_data["correct_response_lengths"]
+            )
+            for p, val in correct_response_length_percentile.items():
+                rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses
+            if len(correct_entropy) > 0:
+                sum_of_sample_mean = get_sum_of_sample_mean(
+                    correct_total_lengths, correct_response_lengths, correct_loss_masks
+                )
+                correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0))
+                rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses
+            else:
+                rollout_data["correct_entropy"] = [0] * num_correct_responses
+
 
 def log_multi_turn_data(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) -> None:
     """

slime/backends/megatron_utils/initialize.py

@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 from megatron.core import mpu, tensor_parallel
+from megatron.core.config import set_experimental_flag
 from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator
 from megatron.training.global_vars import _build_tokenizer, set_args
 
@@ -54,6 +55,10 @@ def _initialize_distributed(args, get_embedding_ranks=None, get_position_embeddi
 
 def init(args):
     set_args(args)
+    if args.enable_experimental:
+        logger.info("Enable megatron experimental")
+        set_experimental_flag(True)
+
     # Pytorch distributed.
     _initialize_distributed(args)
 

slime/rollout/sglang_rollout.py

@@ -214,6 +214,10 @@ async def generate_and_rm(
     sampling_params: dict[str, Any],
     evaluation: bool = False,
 ) -> Sample | list[Sample]:
+    # mask previous off-policy generation for partial rollout
+    if args.partial_rollout and args.mask_offpolicy_in_partial_rollout and sample.response_length > 0:
+        sample.loss_mask = [0] * sample.response_length
+
     # For samples with existing response, check if they're complete
     if sample.status == Sample.Status.COMPLETED or sample.status == Sample.Status.TRUNCATED:
         assert sample.response is not None

slime/utils/arguments.py

@@ -316,6 +316,15 @@ def add_rollout_arguments(parser):
                     "This is useful for long responses."
                 ),
             )
+            parser.add_argument(
+                "--mask-offpolicy-in-partial-rollout",
+                action="store_true",
+                default=False,
+                help=(
+                    "Whether to mask previous generation in partial rollout. "
+                    "If set, only on-policy generated tokens will be used in training"
+                ),
+            )
             parser.add_argument(
                 "--custom-generate-function-path",
                 type=str,
@@ -600,6 +609,12 @@ def add_eval_arguments(parser):
                     "When provided, this overrides --eval-prompt-data."
                 ),
             )
+            parser.add_argument(
+                "--skip-eval-before-train",
+                action="store_true",
+                default=False,
+                help="Whether to skip evaluation before training.",
+            )
 
             # The following keys are used to override the rollout version during eval.
             parser.add_argument("--eval-input-key", type=str, default=None, help="JSON dataset key")
@@ -922,6 +937,12 @@ def add_wandb_arguments(parser):
                     "Specify the key in the reward dict using this argument.",
                 ),
             )
+            parser.add_argument(
+                "--log-correct-samples",
+                action="store_true",
+                default=False,
+                help="Whether to turn on passrate logging, which will log the pass@n of the responses in the rollout.",
+            )
             parser.add_argument("--wandb-run-id", type=str, default=None)
             return parser
 

slime_plugins/models/hf_attention.py

@@ -5,7 +5,6 @@
 from megatron.core import mpu, tensor_parallel
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.module import MegatronModule
 from transformers import AutoConfig
 
@@ -23,7 +22,7 @@ def __init__(
         config,
         layer_number: int,
         cp_comm_type: str = "p2p",
-        pg_collection: ProcessGroupCollection = None,
+        pg_collection=None,
     ):
         super().__init__(config=config)
         self.args = args

slime_plugins/models/qwen3_next.py

@@ -4,7 +4,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
-from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import get_num_layers_to_build
 from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
@@ -170,7 +169,7 @@ def __init__(
         config,
         layer_number: int,
         cp_comm_type: str = "p2p",
-        pg_collection: ProcessGroupCollection = None,
+        pg_collection=None,
     ):
         super().__init__(
             args,

train.py

@@ -62,7 +62,7 @@ def onload_rollout():
     # train loop.
     # note that for async training, one can change the position of the sync operation(ray.get).
     for rollout_id in range(args.start_rollout_id, args.num_rollout):
-        if args.eval_interval is not None and rollout_id == 0:
+        if args.eval_interval is not None and rollout_id == 0 and not args.skip_eval_before_train:
             ray.get(rollout_manager.eval.remote(rollout_id))
 
         rollout_data_ref = ray.get(rollout_manager.generate.remote(rollout_id))