[Train] Micro-batch scheduling on rollout side + variable global batch size (port slime #1926+#1930+#1933)

.github/workflows/pr-test.yml

@@ -465,7 +465,7 @@ jobs:
             '
 
 
-  e2e-test-plugin-contracts:
+  cpu-unittest:
     needs: pre-commit
 
     if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch'
@@ -476,7 +476,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
+        info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_dp_schedule.py"}, {"num_gpus": 0, "test_file": "test_cp_utils.py"}, {"num_gpus": 0, "test_file": "test_metric_report.py"}, {"num_gpus": 0, "test_file": "test_metric_report_dist.py"}, {"num_gpus": 0, "test_file": "test_loss_cp_invariance.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}

.github/workflows/pr-test.yml.j2

@@ -49,12 +49,17 @@
       ],
     },
 
-    'e2e-test-plugin-contracts': {
-      'label': 'run-ci-plugin-contracts',
+    'cpu-unittest': {
+      'label': 'run-ci-cpu-unittest',
       'always': True,
       'cpu': True,
       'tests': [
         {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_dp_schedule.py', 'num_gpus': 0},
+        {'test_file': 'test_cp_utils.py', 'num_gpus': 0},
+        {'test_file': 'test_metric_report.py', 'num_gpus': 0},
+        {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0},
+        {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_path_loading_contracts.py', 'num_gpus': 0},

docs/en/get_started/customization.md

@@ -435,7 +435,7 @@ python -m pytest \
 
 Each test file can also be executed directly with `python tests/plugin_contracts/<file>.py`, which keeps them compatible with `run-ci-changed`.
 
-A dedicated `run-ci-plugin-contracts` CI label is also available. Adding it to a PR triggers all four contract test files in parallel (no GPU required).
+A dedicated `run-ci-cpu-unittest` CI label is also available. Adding it to a PR triggers the CPU-only unit-test job, which runs the contract tests plus other lightweight unit tests in parallel (no GPU required).
 
 For user-defined implementations, you can either export environment variables such as `SLIME_CONTRACT_ROLLOUT_FUNCTION_PATH` and `SLIME_CONTRACT_CUSTOM_RM_PATH`, or pass overrides directly when running a test file, for example:
 

docs/zh/get_started/customization.md

@@ -437,7 +437,7 @@ python -m pytest \
 
 每个测试文件也支持直接通过 `python tests/plugin_contracts/<file>.py` 执行，这样可以和 `run-ci-changed` 保持兼容。
 
-CI 中也提供了独立的 `run-ci-plugin-contracts` label，给 PR 打上该标签后会并行运行上述全部四个契约测试（无需 GPU）。
+CI 中也提供了独立的 `run-ci-cpu-unittest` label，给 PR 打上该标签后会并行运行 CPU-only 的单元测试任务，包含上述契约测试以及其他轻量单测（无需 GPU）。
 
 如果你要验证自己的自定义实现，可以直接设置环境变量，例如 `SLIME_CONTRACT_ROLLOUT_FUNCTION_PATH`、`SLIME_CONTRACT_CUSTOM_RM_PATH`，也可以在直接运行测试文件时传参，例如：
 

examples/multi_agent/agent_system.py

@@ -194,6 +194,19 @@ async def run_agent_system(args, sample):
     args = deepcopy(args)  # Deep copy args because rollout_with_multi_agents mutates them.
     args.sample = sample
     args.results_dict = {"solver": [], "rewriter": [], "selector": []}
+    # Every sample emitted below is a training sample split out of this one
+    # rollout execution (the input ``sample``). Stamp the shared rollout id on
+    # every collected sample at each return point so the per-rollout loss
+    # reducer aggregates the solver / rewriter / selector siblings as one
+    # rollout instead of N, and the by-rollout step splitter keeps them in
+    # the same step. Captured here because ``sample`` gets shadowed by zip-
+    # loop variables further down.
+    input_rollout_id = sample.index
+
+    def _emit(samples_list):
+        for s in samples_list:
+            s.rollout_id = input_rollout_id
+        return samples_list
 
     problem_statement = sample.prompt
     tasks = [solver_worker(args, problem_statement, worker_id) for worker_id in range(args.num_parallel)]
@@ -212,7 +225,7 @@ def reward_adjustment(samples, reward_weight):
 
     if len(previous_solutions) == 0:
         reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
-        return args.results_dict["solver"]
+        return _emit(args.results_dict["solver"])
 
     # Rewriting
     tasks = [
@@ -234,15 +247,15 @@ def reward_adjustment(samples, reward_weight):
     if len(rewrited_solutions) == 0:
         reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
         reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
-        return args.results_dict["solver"] + args.results_dict["rewriter"]
+        return _emit(args.results_dict["solver"] + args.results_dict["rewriter"])
 
     # Selection
     selector = SelectorAgent()
     response = await selector.select(args, problem_statement, rewrited_solutions)
     if len(args.results_dict["selector"]) == 0:
         reward_adjustment(args.results_dict["solver"], args.incorrect_reward_weight)
         reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
-        return args.results_dict["solver"] + args.results_dict["rewriter"]
+        return _emit(args.results_dict["solver"] + args.results_dict["rewriter"])
 
     assert (
         len(args.results_dict["selector"]) == 1
@@ -271,4 +284,4 @@ def reward_adjustment(samples, reward_weight):
         reward_adjustment(args.results_dict["rewriter"], args.incorrect_reward_weight)
         reward_adjustment(args.results_dict["selector"], args.incorrect_reward_weight)
 
-    return args.results_dict["solver"] + args.results_dict["rewriter"] + args.results_dict["selector"]
+    return _emit(args.results_dict["solver"] + args.results_dict["rewriter"] + args.results_dict["selector"])

tests/_cp_dist_helpers.py

@@ -0,0 +1,167 @@
+"""Shared infrastructure for the CP-related multi-process CPU tests.
+
+Why this module exists
+----------------------
+The CP / metric-report / backward-grad-norm tests all want to:
+
+1. Stub ``megatron.core.mpu`` *before* importing
+   ``vime.backends.megatron_utils.cp_utils`` (the CPU CI image has no real
+   megatron).
+2. Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` and
+   exercise the actual production helpers (``get_sum_of_sample_mean``,
+   ``reduce_train_step_metrics``, ``gather_and_reduce_log_dict``,
+   ``rollout_log_metric_contribution``).
+3. Chunk each sample's response tensor across CP ranks the same way the
+   real forward pass does — using
+   ``get_logits_and_tokens_offset_with_cp`` so the slicing stays in lock-
+   step with the production reducer.
+
+Putting that here keeps the per-feature test files focused on the
+behaviour they check (numerics / report formulas / backward) rather than
+on plumbing.
+
+Mapping to Megatron
+-------------------
+- ``mp.spawn(...)`` + gloo backend mirrors the per-rank entry-point that
+  ``torch.distributed.run`` would create for a real launch.
+- ``dp_cp_group = new_group(range(world_size))`` matches
+  ``parallel_state.get_data_parallel_group(with_context_parallel=True)``
+  (Megatron-LM ``finalize_model_grads.py:437``).  In the no-TP / no-PP
+  CPU test setup the whole world *is* that group.
+- The per-rank CP chunking mirrors what the attention layer feeds into
+  the loss in Megatron: each CP rank only sees its 2-chunk slice of the
+  response tokens (cf. ``cp_utils.get_logits_and_tokens_offset_with_cp``,
+  the same helper used by the real forward pass).
+"""
+
+from __future__ import annotations
+
+import os
+import socket
+import sys
+import types
+
+
+# --- Stub ``megatron.core.mpu`` (must run before cp_utils is imported) ---
+#
+# Both this module and any test file that imports it should *import this
+# helper first*. Doing so installs the stub at import time so that the
+# subsequent ``from vime.backends.megatron_utils.cp_utils import ...`` in
+# the test file binds ``cp_utils.mpu`` to this stub.
+#
+# In spawned workers, ``mp.spawn`` re-imports the test module fresh, which
+# re-runs this stub installation; then the worker mutates the stub's
+# ``get_context_parallel_*`` attributes via ``_stub_megatron_in_worker``
+# below to pin (cp_size, cp_rank) for that worker.
+_fake_mpu = types.ModuleType("megatron.core.mpu")
+_fake_mpu.get_context_parallel_world_size = lambda: 1
+_fake_mpu.get_context_parallel_rank = lambda: 0
+_fake_core = types.ModuleType("megatron.core")
+_fake_core.mpu = _fake_mpu
+_fake_megatron = types.ModuleType("megatron")
+_fake_megatron.core = _fake_core
+sys.modules.setdefault("megatron", _fake_megatron)
+sys.modules.setdefault("megatron.core", _fake_core)
+sys.modules.setdefault("megatron.core.mpu", _fake_mpu)
+
+
+def stub_megatron_in_worker(cp_size: int, cp_rank: int) -> None:
+    """Override ``mpu.get_context_parallel_*`` inside an ``mp.spawn`` worker.
+
+    ``mp.spawn`` pickles the worker function by name and re-imports the
+    test module in the child — that re-runs the top-of-file stub install
+    with ``cp_size=1``. By the time the worker runs, ``cp_utils`` has
+    already bound its module-level ``mpu`` reference to the stub.
+
+    So we must MUTATE the stub module's attributes in place rather than
+    replace ``sys.modules['megatron.core.mpu']`` — replacing the module
+    would leave ``cp_utils.mpu`` pointing at the now-shadowed stub.
+    """
+    from megatron.core import mpu  # the stub installed at import time
+
+    mpu.get_context_parallel_world_size = lambda: cp_size
+    mpu.get_context_parallel_rank = lambda: cp_rank
+
+
+def free_port() -> int:
+    """Pick an unused TCP port for ``init_process_group``'s rendezvous.
+
+    Equivalent to what ``torchrun`` does when ``--master-port`` is not
+    set; we just need a port nothing else is bound to so multiple
+    parametrized test cases can spawn without colliding.
+    """
+    s = socket.socket()
+    s.bind(("", 0))
+    port = s.getsockname()[1]
+    s.close()
+    return port
+
+
+def init_worker_process_group(rank: int, world_size: int, master_port: int):
+    """Stand up gloo ``torch.distributed`` and return the DP*CP group.
+
+    The CPU CI image ships gloo but not NCCL; in the no-TP / no-PP setup
+    the DP-with-CP group is the whole world, mirroring
+    ``parallel_state.get_data_parallel_group(with_context_parallel=True)``
+    in Megatron-LM ``finalize_model_grads.py:437``.
+    """
+    import torch.distributed as _dist
+
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(master_port)
+    _dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
+    return _dist.new_group(ranks=list(range(world_size)))
+
+
+def cp_chunk_response_tensor(x, total_length: int, response_length: int):
+    """Slice a sample's response tensor to what the current CP rank sees.
+
+    Mirrors the real forward pass: at CP > 1 each rank's attention only
+    consumes the two response-token chunks selected by
+    ``get_logits_and_tokens_offset_with_cp`` (the same helper used by the
+    production reducer in ``cp_utils.get_sum_of_sample_mean``). So the
+    "x" we feed into the reducer on a CP rank must be sliced the same
+    way to keep the numbers honest.
+
+    Importing locally so callers don't pay the import cost before
+    ``stub_megatron_in_worker`` has had a chance to pin (cp_size, cp_rank).
+    """
+    import torch
+
+    from vime.backends.megatron_utils.cp_utils import get_logits_and_tokens_offset_with_cp
+
+    prompt_length = total_length - response_length
+    _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length)
+    c0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length]
+    c1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length]
+    return torch.cat([c0, c1])
+
+
+# ---------------------------------------------------------------------------
+# Shared four-rollout fixture, used by both the metric-report distributed
+# tests and the backward-grad-norm test. Keeping the data in one place so
+# the "train report matches rollout report matches grad-norm baseline"
+# contract is anchored on the same numbers everywhere.
+#
+# Four samples (1 rollout each), total_length=12 (4 prompt + 8 response),
+# loss_mask=all-ones. x values differ by orders of magnitude so any cross-
+# rank summation bug shows up as a visibly wrong number.
+#
+# Per-sample token-mean: 4.5 / 45 / 450 / 4500.
+# Per-rollout-mean report (sum / num_rollouts):
+#       (4.5 + 45 + 450 + 4500) / 4 = 1249.875
+# Per-token-loss report (sum_x / total_tokens):
+#       (36 + 360 + 3600 + 36000) / 32 = 1249.875
+# (the two paths agree by construction so the test expectations stay
+# simple — the *report formulas* are still distinct as exercised inside
+# ``reduce_train_step_metrics``.)
+# ---------------------------------------------------------------------------
+FOUR_ROLLOUT_TOTAL_LENGTHS = [12, 12, 12, 12]
+FOUR_ROLLOUT_RESPONSE_LENGTHS = [8, 8, 8, 8]
+FOUR_ROLLOUT_X_VALUES = [
+    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
+    [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0],
+    [100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0],
+    [1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0, 7000.0, 8000.0],
+]
+FOUR_ROLLOUT_EXPECTED_REPORT = 1249.875