From de09f3b14bf82ea748eea1dfa0ba2b23bf37d439 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 2 Jun 2026 04:40:52 +0000 Subject: [PATCH] [Train] Move micro-batch scheduling to rollout side + variable global batch size (port slime #1926+#1930+#1933) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the coupled THUDM/slime chain (one feature line, three sequential PRs) into vime as a single change since they touch the same files interdependently: - #1926: move the per-step micro-batch schedule from the actor (train) side to `_split_train_data_by_dp` on the rollout side, computed once over the full global batch. New `vime/utils/dp_schedule.py` (`build_dp_schedule`, `compute_dynamic_global_batch_size`) holds the Ray/engine-free logic so it is unit-testable. Drops the DP-wide `all_reduce(num_microbatches, MAX)`; `get_data_iterator` becomes a thin consumer reading `num_microbatches` / `micro_batch_indices` straight out of `rollout_data`. - #1930 + #1933 [1/N,2/N]: variable global batch size — per-step split by rollout id so each rollout yields a fixed training-step count regardless of how many samples it produced; CP-aware loss/metric handling. vime reconciliation notes: - vime already carried `--use-dynamic-global-batch-size` / `--disable-rollout-trim-samples` (args + inline `_compute_dynamic_global_batch_size`) from an earlier partial sync; #1926 removes that inline method in favor of `dp_schedule.compute_dynamic_global_batch_size`, applied cleanly. - `_split_train_data_by_dp` keeps vime's diverged extra keys (`rollout_routed_experts`, `teacher_log_probs`, `dynamic_global_batch_size`). - `train_parallel_config` is extended with `cp_size` / `vpp_size` / `microbatch_group_size_per_vp_stage` (set in actor.py) and read by `build_dp_schedule`. - import-block hunks hand-reconciled (vime's vllm seam, no SGLangEngine import). Adds the chain's CPU tests (registered in the always-on cpu job): test_dp_schedule, test_cp_utils, test_loss_cp_invariance, test_metric_report, test_metric_report_dist (+ tests/_cp_dist_helpers.py). NOTE: this changes training micro-batch counts / bin-packing (numerics) — needs GPU validation before merge; cpu validation here covers the shipped CPU tests. Signed-off-by: aoshen02 --- .github/workflows/pr-test.yml | 4 +- .github/workflows/pr-test.yml.j2 | 9 +- docs/en/get_started/customization.md | 2 +- docs/zh/get_started/customization.md | 2 +- examples/multi_agent/agent_system.py | 21 +- tests/_cp_dist_helpers.py | 167 ++++++++++++ tests/test_cp_utils.py | 177 +++++++++++++ tests/test_dp_schedule.py | 288 ++++++++++++++++++++ tests/test_loss_cp_invariance.py | 253 ++++++++++++++++++ tests/test_metric_report.py | 320 +++++++++++++++++++++++ tests/test_metric_report_dist.py | 303 +++++++++++++++++++++ vime/backends/megatron_utils/actor.py | 33 ++- vime/backends/megatron_utils/cp_utils.py | 136 +++++++++- vime/backends/megatron_utils/data.py | 273 +++++++------------ vime/backends/megatron_utils/loss.py | 33 ++- vime/backends/megatron_utils/model.py | 70 +++-- vime/ray/rollout.py | 166 +++++++----- vime/utils/arguments.py | 12 - vime/utils/data.py | 14 - vime/utils/dp_schedule.py | 191 ++++++++++++++ vime/utils/seqlen_balancing.py | 52 ++++ vime/utils/types.py | 8 + 22 files changed, 2205 insertions(+), 329 deletions(-) create mode 100644 tests/_cp_dist_helpers.py create mode 100644 tests/test_cp_utils.py create mode 100644 tests/test_dp_schedule.py create mode 100644 tests/test_loss_cp_invariance.py create mode 100644 tests/test_metric_report.py create mode 100644 tests/test_metric_report_dist.py create mode 100644 vime/utils/dp_schedule.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index bc8b8e769..de5ca00b1 100644 --- a/.github/workflows/pr-test.yml +++ b/.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 }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index adab45a63..853a779d3 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.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}, diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md index c913958f8..e5f2a9b8f 100644 --- a/docs/en/get_started/customization.md +++ b/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/.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: diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md index b885ca58b..76a7d777a 100644 --- a/docs/zh/get_started/customization.md +++ b/docs/zh/get_started/customization.md @@ -437,7 +437,7 @@ python -m pytest \ 每个测试文件也支持直接通过 `python tests/plugin_contracts/.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`，也可以在直接运行测试文件时传参，例如： diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py index 8919b4c23..15c6a4f57 100644 --- a/examples/multi_agent/agent_system.py +++ b/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,7 +247,7 @@ 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() @@ -242,7 +255,7 @@ def reward_adjustment(samples, reward_weight): 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"]) diff --git a/tests/_cp_dist_helpers.py b/tests/_cp_dist_helpers.py new file mode 100644 index 000000000..1382094fe --- /dev/null +++ b/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 diff --git a/tests/test_cp_utils.py b/tests/test_cp_utils.py new file mode 100644 index 000000000..d37e870a7 --- /dev/null +++ b/tests/test_cp_utils.py @@ -0,0 +1,177 @@ +"""CPU unit tests for ``vime.backends.megatron_utils.cp_utils.get_sum_of_sample_mean``. + +Pins the per-rollout reducer contract: a rollout split into N training +samples (compact / subagent) must contribute exactly one token-weighted +mean to the sum, even when first-fit packing puts those siblings into +different micro-batches at training time. + +The CPU-only CI image does not ship megatron — ``_cp_dist_helpers`` +stubs ``megatron.core.mpu`` at import time so the subsequent +``cp_utils`` import binds against the stub. + +End-to-end report-formula invariance and multi-process distributed +checks live in ``test_metric_report.py`` and ``test_metric_report_dist.py``. +""" + +from __future__ import annotations + +# Import the helpers BEFORE the slime imports so the megatron stub lands +# in sys.modules first. pytest's prepend importmode puts this file's +# directory (``tests/``) on sys.path, which is what makes the bare-name +# import work without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from vime.backends.megatron_utils.cp_utils import ( # noqa: E402 + get_logits_and_tokens_offset_with_cp, + get_sum_of_sample_mean, +) + + +def _make_inputs(per_sample_lengths: list[int]): + """Build (total_lengths, response_lengths, loss_masks) for samples of the given lengths. + + Each sample has loss_mask = all-ones (so mask sum == length); total length + is response length + 4 fake prompt tokens (unused by the reducer in + cp_size==1 mode). + """ + response_lengths = list(per_sample_lengths) + total_lengths = [r + 4 for r in response_lengths] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + return total_lengths, response_lengths, loss_masks + + +def _denoms(*values: int) -> torch.Tensor: + """Wrap per-sample denoms as the float tensor that the actor side promotes + them to before calling the reducer.""" + return torch.tensor(values, dtype=torch.float32) + + +@pytest.mark.unit +def test_default_reduces_to_per_sample_mean(): + """``sample_denoms=None`` reproduces the legacy per-sample-mean.""" + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3]) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]) + # per-sample means: 2, 5, 8 → sum = 15 + assert reducer(x).item() == pytest.approx(15.0) + + +@pytest.mark.unit +def test_per_rollout_denom_collapses_siblings_into_one_mean(): + """Pre-computed per-rollout mask sums make N sibling samples contribute one + token-weighted mean instead of N per-sample means.""" + # 4 samples: rollout R0 owns indices 0,1,2 (mask sums 3+3+3=9); rollout R1 + # owns index 3 (mask sum 3). Pre-computed per-sample denom = group sum. + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + sample_denoms = _denoms(9, 9, 9, 3) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + # R0 token-mean: (1+2+...+9)/9 = 5. R1 token-mean: (10+11+12)/3 = 11. Sum = 16. + assert reducer(x).item() == pytest.approx(16.0) + + +@pytest.mark.unit +def test_split_across_mbs_recovers_full_per_rollout_mean(): + """The critical contract: when a rollout's samples land in different mbs, + summing each mb's reducer output equals one whole-step reducer call with + the same pre-computed denominators. This is exactly the bug that motivated + the precomputation — if the denom were computed per-mb (partial mask sum), + the two halves wouldn't add up.""" + # 4 samples (same as above). Whole-step denoms = [9, 9, 9, 3]. + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + sample_denoms = _denoms(9, 9, 9, 3) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + + whole = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + whole_value = whole(x).item() + + # mb_a holds samples 0, 1 of R0; mb_b holds sample 2 of R0 and sample 3 (R1). + # Each mb carries the SAME per-sample denoms (precomputed at step level) + # — that's what makes the split safe. + mb_a = get_sum_of_sample_mean(total_lengths[:2], response_lengths[:2], loss_masks[:2], sample_denoms[:2]) + mb_b = get_sum_of_sample_mean(total_lengths[2:], response_lengths[2:], loss_masks[2:], sample_denoms[2:]) + split_value = mb_a(x[:6]).item() + mb_b(x[6:]).item() + + assert split_value == pytest.approx(whole_value) + + +@pytest.mark.unit +def test_split_with_per_mb_denom_would_be_wrong(): + """Sanity-check the bug we're guarding against: if the caller naively + computes per-rollout denoms from each mb's own samples (the local mask + sum, NOT the precomputed whole-rollout sum), the two halves DON'T add up + to the whole-step value. This pins down WHY the precomputation must + happen at the step level.""" + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + + whole = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, _denoms(9, 9, 9, 3)) + whole_value = whole(x).item() + + # Wrong denom: each mb only sees its own samples of R0. + # mb_a's "rollout mask sum" for R0 would be 3+3=6 (instead of 9). mb_b's + # would be 3. Different from the true whole-rollout total. + mb_a_wrong = get_sum_of_sample_mean(total_lengths[:2], response_lengths[:2], loss_masks[:2], _denoms(6, 6)) + mb_b_wrong = get_sum_of_sample_mean(total_lengths[2:], response_lengths[2:], loss_masks[2:], _denoms(3, 3)) + wrong_total = mb_a_wrong(x[:6]).item() + mb_b_wrong(x[6:]).item() + + assert wrong_total != pytest.approx(whole_value), ( + "Expected the per-mb denom path to produce a different (incorrect) value; " + "if these match, the regression test is no longer guarding the precomputation contract." + ) + + +@pytest.mark.unit +def test_cp_chunking_preserves_per_rollout_mean_report(monkeypatch): + """Turning CP on must not change the reducer's output. + + Real flow: each CP rank only sees its chunk of the response tokens; the + reducer's CP>1 branch slices ``loss_mask`` to match. Summing each CP + rank's reducer output across CP ranks reproduces the cp=1 result, which + is what train_one_step then divides by ``step_global_batch_size``. + """ + from megatron.core import mpu as _mpu + + # Use lengths that line up cleanly with the CP chunking + # (chunk_size = ceil(total_length / (2*cp_size))). + total_lengths = [12, 12] # 2 samples + response_lengths = [8, 8] # 4 prompt + 8 response each + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + sample_denoms = torch.tensor([16.0, 16.0], dtype=torch.float32) # = sum of both mask totals (one rollout) + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + x_concat = torch.cat(x_full) + + # --- cp=1 baseline --- + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: 1) + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda: 0) + reducer_cp1 = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + baseline = reducer_cp1(x_concat).item() + + # --- cp=2: sum partial reducer outputs across the two CP ranks --- + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: 2) + cp_total = 0.0 + for cp_rank in range(2): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + # Slice each sample's response-token tensor to the chunks this CP + # rank owns, mirroring what the forward pass would feed in. + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + chunk_0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + chunk_1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + x_chunks_per_sample.append(torch.cat([chunk_0, chunk_1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer_cp2 = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + cp_total += reducer_cp2(x_for_rank).item() + + assert cp_total == pytest.approx(baseline) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_dp_schedule.py b/tests/test_dp_schedule.py new file mode 100644 index 000000000..6ce03ce2b --- /dev/null +++ b/tests/test_dp_schedule.py @@ -0,0 +1,288 @@ +"""CPU unit tests for vime.utils.dp_schedule.build_dp_schedule. + +The tests assert the invariants documented at the top of dp_schedule.py against +a range of static / dynamic / VPP / oversize / balance / uneven scenarios. +""" + +from types import SimpleNamespace + +import pytest + +from vime.utils.dp_schedule import build_dp_schedule + + +def make_args( + *, + micro_batch_size=1, + use_dynamic_batch_size=False, + max_tokens_per_gpu=None, + balance_data=False, +): + return SimpleNamespace( + micro_batch_size=micro_batch_size, + use_dynamic_batch_size=use_dynamic_batch_size, + max_tokens_per_gpu=max_tokens_per_gpu, + balance_data=balance_data, + ) + + +def make_tp(dp_size=1, cp_size=1, vpp_size=1, microbatch_group_size_per_vp_stage=1): + return { + "dp_size": dp_size, + "cp_size": cp_size, + "vpp_size": vpp_size, + "microbatch_group_size_per_vp_stage": microbatch_group_size_per_vp_stage, + } + + +def assert_invariants( + partitions, + micro_batch_indices, + num_microbatches, + *, + dp_size, + expected_global_sample_indices, + total_lengths, + max_per_bin=None, +): + """Check the invariants documented at the top of dp_schedule.py. + + ``expected_global_sample_indices`` is the set of global sample indices + that should end up covered (after trim). Trailing rollouts that don't + fit are excluded. + """ + seen_global: set[int] = set() + for r in range(dp_size): + partition = partitions[r] + mbi = micro_batch_indices[r] + + # Same num_mbs per rank (PP sync). + assert len(mbi) == sum(num_microbatches), f"rank {r}: mbs count mismatch" + + # Flattened micro_batch_indices == range(len(partition)). + flat = [i for mbs in mbi for i in mbs] + assert flat == list(range(len(partition))), f"rank {r}: micro_batch_indices don't tile [0, n)" + + # Disjoint partitions whose union covers every kept sample. + assert seen_global.isdisjoint(partition), f"rank {r}: overlap with other ranks" + seen_global.update(partition) + assert seen_global == set(expected_global_sample_indices), "covered sample set mismatch" + + if max_per_bin is None: + return + + # Every mbs <= max_per_bin tokens, EXCEPT a singleton bin holding an oversized sample. + for r in range(dp_size): + partition = partitions[r] + for mbs in micro_batch_indices[r]: + bin_total = sum(total_lengths[partition[i]] for i in mbs) + if bin_total > max_per_bin: + assert len(mbs) == 1, f"rank {r}: mbs sum {bin_total} > {max_per_bin} but contains {len(mbs)} samples" + + +@pytest.mark.unit +def test_static_stride_single_step(): + """Static + strided DP split, single step (1 rollout = 1 sample).""" + total_lengths = [10] * 16 + rollout_indices = list(range(16)) + args = make_args(micro_batch_size=2) + tp = make_tp(dp_size=4) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=16, rollout_indices=rollout_indices + ) + + assert nmb == [2] + assert gbs_per_step == [16] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=4, + expected_global_sample_indices=range(16), + total_lengths=total_lengths, + ) + + +@pytest.mark.unit +def test_static_balance_multi_step(): + """Static + balance_data + 2 training steps.""" + total_lengths = [1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1] + rollout_indices = list(range(16)) + args = make_args(micro_batch_size=2, balance_data=True) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert nmb == [2, 2] + assert gbs_per_step == [8, 8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(16), + total_lengths=total_lengths, + ) + + +@pytest.mark.unit +def test_dynamic_uniform(): + """Dynamic mbs on uniform-length samples.""" + total_lengths = [5] * 8 + rollout_indices = list(range(8)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=10, + ) + + +@pytest.mark.unit +def test_dynamic_oversized_sample_lands_alone(): + """A sample larger than max_per_bin must end up alone in its mbs.""" + total_lengths = [15, 3, 3, 3, 3, 3, 3, 3] + rollout_indices = list(range(8)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=10, + ) + oversize_idx = total_lengths.index(15) + found = False + for r in range(2): + if oversize_idx not in partitions[r]: + continue + local = partitions[r].index(oversize_idx) + for mbs in mbi[r]: + if local in mbs: + assert mbs == [local], f"oversized sample shares an mbs: {mbs}" + found = True + assert found + + +@pytest.mark.unit +def test_dynamic_with_vpp_rounds_to_mb_group(): + """num_microbatches per rank should be a multiple of mb_group when vpp_size > 1.""" + total_lengths = [4] * 32 + rollout_indices = list(range(32)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=8) + tp = make_tp(dp_size=2, vpp_size=2, microbatch_group_size_per_vp_stage=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=16, rollout_indices=rollout_indices + ) + + for n in nmb: + assert n % 2 == 0, f"num_microbatches {n} is not a multiple of mb_group=2" + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(32), + total_lengths=total_lengths, + max_per_bin=8, + ) + + +@pytest.mark.unit +def test_rollout_grouping_keeps_samples_together(): + """compact / subagent simulation: rollout 0 emits 3 samples, rollout 1 emits 2, + rollout 2 emits 4. Splitter keeps every rollout's samples in a single step.""" + rollout_indices = [0, 0, 0, 1, 1, 2, 2, 2, 2] + total_lengths = [3] * 9 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=1, rollout_indices=rollout_indices + ) + + # 3 rollouts / 1 per step → 3 steps, gbs constant. + assert gbs_per_step == [1, 1, 1] + # For each step, collect the samples (global indices) that landed in that step's mbs + # on rank 0, then verify they exactly equal the rollout's sample positions. + expected_per_step = [[0, 1, 2], [3, 4], [5, 6, 7, 8]] + rank0_partition = partitions[0] + mbs_cursor = 0 + for step_i, n_mbs in enumerate(nmb): + step_locals = sorted(j for mbs in mbi[0][mbs_cursor : mbs_cursor + n_mbs] for j in mbs) + step_globals = [rank0_partition[j] for j in step_locals] + assert ( + sorted(step_globals) == expected_per_step[step_i] + ), f"step {step_i} samples = {step_globals}, expected {expected_per_step[step_i]}" + mbs_cursor += n_mbs + assert_invariants( + partitions, + mbi, + nmb, + dp_size=1, + expected_global_sample_indices=range(9), + total_lengths=total_lengths, + max_per_bin=12, + ) + + +@pytest.mark.unit +def test_trims_trailing_rollouts_that_dont_fill_a_step(): + """5 rollouts, gbs=2 → 2 steps × 2 rollouts; trailing rollout 4 (sample positions 6, 7) + is dropped.""" + rollout_indices = [0, 0, 1, 2, 2, 3, 4, 4] + total_lengths = [3] * 8 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=2, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [2, 2] + # Sample positions 6 and 7 belong to the trimmed rollout 4 and must be absent. + assert_invariants( + partitions, + mbi, + nmb, + dp_size=1, + expected_global_sample_indices=range(6), + total_lengths=total_lengths, + max_per_bin=12, + ) + + +@pytest.mark.unit +def test_rejects_when_fewer_rollouts_than_gbs(): + """gbs=4 with only 3 distinct rollouts → cannot form one step.""" + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + with pytest.raises(AssertionError, match="num_rollouts"): + build_dp_schedule(args, tp, [3] * 6, global_batch_size=4, rollout_indices=[0, 0, 1, 1, 2, 2]) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_loss_cp_invariance.py b/tests/test_loss_cp_invariance.py new file mode 100644 index 000000000..998ba1d6c --- /dev/null +++ b/tests/test_loss_cp_invariance.py @@ -0,0 +1,253 @@ +"""End-to-end backward / gradient-norm CP-invariance check on CPU. + +This is the closest thing to a real training-step backward we can run on +the CPU CI image without standing up Megatron, FlashAttention, or NCCL. +The goal: prove that for the same training samples, *the gradient norm +after the optimizer-side all-reduce is identical regardless of CP size*. + +Why this matters +---------------- +Slime's loss prescaling + Megatron's per-mb scaling + DDP's grad +averaging compose into one big formula. Any time we touch any one of +those three layers the numbers should land in the same place. Until +this test existed we only had end-to-end report-formula checks +(`test_metric_report_dist.py`); none of them ran a real ``backward()``, +so a sign or factor error in the prescaling would slip through. + +Mapping to Megatron source +-------------------------- +We reproduce, for each spawned rank, the exact sequence Megatron applies +when a 3-tuple ``(loss, num_tokens, log)`` comes back from the loss +function with ``calculate_per_token_loss=False`` — slime's per-rollout- +mean path: + + 1. Loss function pre-scales:: + loss *= num_microbatches / step_global_batch_size * (dp * cp) + See ``slime/backends/megatron_utils/loss.py:1209-1215``. + 2. Megatron divides by ``clamp(num_tokens, 1)`` then by + ``num_microbatches``:: + output_tensor /= torch.clamp(num_tokens, min=1) # num_tokens=1 → no-op + output_tensor /= num_microbatches + See ``Megatron-LM/megatron/core/pipeline_parallel/schedules.py:258-264`` + (the ``len(outputs) == 3`` branch with ``not calculate_per_token_loss``). + 3. Backward fills grad buffers; per-mb contributions sum on each rank. + 4. DDP grad sync averages across the DP-with-CP group:: + grad_sum_across_dp_cp_world / (dp * cp) + See ``Megatron-LM/megatron/core/distributed/distributed_data_parallel.py:283-290`` + (``average_in_collective=False``, ``gradient_scaling_factor = 1.0 / dp_cp_group.size()``). + +Composing 1-4 collapses to + final_grad = total_sum_of_rollout_means / step_global_batch_size, +i.e. the gradient of ``mean_of_per_rollout_means(x)``. That doesn't +contain ``cp`` anywhere, so the grad norm must be identical for any +(dp, cp) factorization of the same world size. + +What this test does NOT exercise: the actual Megatron model classes, the +real DDP buffer code, fused optimizers, mixed-precision. We use a plain +``nn.Linear`` with manual all-reduce-average to simulate steps 1-4 above. +The contract here is on *our* scaling math (steps 1 + 4 are slime's; +step 2 is what Megatron does to our 3-tuple). If Megatron later changes +step 2 — e.g. drops the ``/= num_microbatches`` — this test won't catch +it, but the real GPU integration suite (``test_qwen2.5_0.5B_short.py``) +will. +""" + +from __future__ import annotations + +# Megatron stub must land in sys.modules first; the slime imports inside +# the worker pick it up via this same module. pytest's prepend importmode +# puts ``tests/`` on sys.path so the bare-name import works without an +# ``__init__.py``; mp.spawn children inherit the parent's sys.path. +import _cp_dist_helpers +import pytest +import torch +from _cp_dist_helpers import ( + FOUR_ROLLOUT_EXPECTED_REPORT, + FOUR_ROLLOUT_RESPONSE_LENGTHS, + FOUR_ROLLOUT_TOTAL_LENGTHS, + FOUR_ROLLOUT_X_VALUES, + cp_chunk_response_tensor, + free_port, + init_worker_process_group, + stub_megatron_in_worker, +) + + +def _grad_norm_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + seed: int, + master_port: int, + result_path: str, +) -> None: + """One spawned rank. + + Builds a tiny ``nn.Linear`` model (deterministic init via ``seed``), + runs slime's per-rollout-mean loss reducer with the rank's share of + the four-rollout fixture, applies the slime-side prescaling, then + Megatron's per-mb scaling, then ``.backward()``, then a manual + all-reduce-average across the dp-with-cp group (mirroring DDP's + ``average_in_collective=False`` path with + ``gradient_scaling_factor = 1 / dp_cp_world_size``). Rank 0 writes the + final ``grad_norm`` to ``result_path``. + """ + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_cp_group = init_worker_process_group(rank, world_size, master_port) + try: + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean + + # Same init across all (dp, cp) configs so the grad we backprop + # into is comparable. ``manual_seed`` is enough on CPU because we + # only do one forward/backward and no dropout. + torch.manual_seed(seed) + model = torch.nn.Linear(1, 1, bias=False) + # Force a known weight value to keep the math hand-checkable: with + # weight = 1.0 and input = x, the linear output equals x, and the + # grad of (output * x).sum() wrt weight equals (x*x).sum(). That + # makes the gradient a pure function of the fixture's x values, + # independent of the random init draw. + with torch.no_grad(): + model.weight.fill_(1.0) + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + step_global_batch_size = 4 # 4 rollouts in the step + num_microbatches = 1 # this CPU model does the whole rank-share in one mb + + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x_full = [all_x[i] for i in my_indices] + # Pre-computed per-rollout denoms = each sample's own mask sum + # (each rollout in the fixture has exactly one sample, so the + # per-rollout denom collapses to the per-sample denom). + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x_full) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x_full, strict=True)] + ) + + # === Forward path ===================================================== + # Tiny "model": output[i] = x[i] * weight. We treat the linear + # output as the per-token quantity the loss is computed over — + # this stands in for the (logits @ token_emb) the policy loss + # consumes in real training. + x_input = x_for_rank.unsqueeze(-1) # shape [T, 1] + output = model(x_input).squeeze(-1) # shape [T] + + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + loss = reducer(output) + + # === Step 1: slime's per-rollout-mean prescaling ====================== + # loss.py:1209-1215. ``mpu.get_data_parallel_world_size(with_context_parallel=True)`` + # is the dp-with-cp world size, which is ``world_size`` in this setup. + loss = loss * num_microbatches / step_global_batch_size * world_size + + # === Step 2: Megatron's forward_step_calc_loss scaling ================ + # schedules.py:258-264 — for the 3-tuple, not-per-token-loss path: + # output_tensor /= torch.clamp(num_tokens, min=1) + # output_tensor /= num_microbatches + # slime passes num_tokens=1 in this path (loss.py:1221), so the + # first divide is a no-op; we keep it explicit to mirror the + # source faithfully. + num_tokens_for_scaling = torch.tensor(1.0) # slime's placeholder + loss = loss / torch.clamp(num_tokens_for_scaling, min=1.0) + loss = loss / num_microbatches + + # === Step 3: backward fills per-rank grad ============================= + loss.backward() + + # === Step 4: DDP all-reduce-average across dp-with-cp world =========== + # distributed_data_parallel.py:283-290, ``average_in_collective=False`` + # case: ``gradient_scaling_factor = 1.0 / dp_cp_group.size()`` is + # baked into the buffer, so the all-reduce is a SUM and the + # 1/world_size scaling pre-applies. We do the equivalent here by + # all-reducing then dividing. + grad = model.weight.grad.detach() + _dist.all_reduce(grad, group=dp_cp_group) + grad = grad / world_size + + # The norm of a 1-element gradient is its absolute value. We + # report ``grad.item()`` directly so the assertion side can also + # eyeball the sign, which is more useful than a strict norm when + # debugging a regression. + grad_value = grad.item() + + if rank == 0: + with open(result_path, "w") as f: + f.write(repr(grad_value)) + finally: + _dist.destroy_process_group() + + +def _run_grad_norm_worker(dp_size: int, cp_size: int, tmp_path) -> float: + """Spawn ``dp_size * cp_size`` workers and return rank-0's final grad.""" + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / f"grad_dp{dp_size}_cp{cp_size}.txt") + mp.spawn( + _grad_norm_worker, + args=(world_size, cp_size, dp_size, 0, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + return float(f.read()) + + +# Subset of (dp, cp) configs to keep runtime down; covers the four +# qualitatively distinct cases: +# - (1, 1) baseline (no parallelism) +# - (2, 1) DP-only +# - (1, 2) CP-only +# - (2, 2) DP + CP combined +# - (1, 4) deeper CP-only +# - (4, 1) deeper DP-only +# The full 3*3 matrix lives in test_metric_report_dist.py — here we just +# want enough coverage to catch a sign/factor regression in the slime +# prescaling math. +_PARALLELISM_CASES = [(1, 1), (2, 1), (1, 2), (2, 2), (1, 4), (4, 1)] + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_size,cp_size", _PARALLELISM_CASES) +def test_backward_grad_is_cp_invariant(dp_size, cp_size, tmp_path): + """The post-DDP-average gradient must be identical across all + (dp, cp) configurations of the same global batch. + + Hand-derivable expectation: with weight = 1.0 and the fixture above, + the gradient of ``mean_of_per_rollout_means(model(x))`` wrt weight is + the same quantity the rollout-report tests pin + (FOUR_ROLLOUT_EXPECTED_REPORT = 1249.875), because for each rollout + the per-token mean of ``x * weight`` differentiates to the per-token + mean of ``x``. + """ + grad = _run_grad_norm_worker(dp_size=dp_size, cp_size=cp_size, tmp_path=tmp_path) + # Tolerance: float32 / multi-rank sums introduce ~1e-3 relative error + # on numbers up to ~1250; that's still 5+ digits of agreement. Each + # (dp, cp) case is pinned to the same hand-derived value, so a sign + # or factor regression in the prescaling math will fail the whole + # matrix uniformly — easy to spot in CI logs. + assert grad == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT, rel=1e-4) + + +# Keep the helpers import load-bearing (it installs the megatron stub). +_ = _cp_dist_helpers + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_metric_report.py b/tests/test_metric_report.py new file mode 100644 index 000000000..98d7d69c6 --- /dev/null +++ b/tests/test_metric_report.py @@ -0,0 +1,320 @@ +"""Single-process metric-report invariance tests. + +Pins train-side / rollout-side report formulas implemented in +``vime.backends.megatron_utils.cp_utils.reduce_train_step_metrics`` and +``rollout_log_metric_contribution``: the reported number for a given set +of samples must be the same regardless of + + - how samples are distributed across micro-batches / DP ranks + - whether context parallelism is on or off + - whether the path is per-rollout-mean or per-token-loss + +Single-process variants use a mock dp-with-cp group + a no-op +``dist.all_reduce`` to keep things lightweight; the multi-process +end-to-end variants (real torch.distributed) live in +``test_metric_report_dist.py``. +""" + +from __future__ import annotations + +# Import the helpers BEFORE the slime imports so the megatron stub lands +# in sys.modules first. pytest's prepend importmode puts this file's +# directory (``tests/``) on sys.path, which is what makes the bare-name +# import work without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from vime.backends.megatron_utils.cp_utils import ( # noqa: E402 + get_logits_and_tokens_offset_with_cp, + get_sum_of_sample_mean, + reduce_train_step_metrics, + rollout_log_metric_contribution, +) + + +@pytest.fixture +def mock_dp_with_cp_group(monkeypatch): + """A sentinel "process group" object plus a no-op ``dist.all_reduce``. + + Lets the train-step report tests exercise the production call shape + (``dist.all_reduce(values, group=dp_with_cp_group)``) without standing + up a real torch.distributed runtime. The test itself simulates cross- + rank summation in pure Python and feeds the already-summed tensor in; + the no-op all_reduce leaves it untouched. + """ + import torch.distributed as dist + + monkeypatch.setattr(dist, "all_reduce", lambda tensor, group=None, op=None: None) + return object() # opaque sentinel — only used as the ``group`` argument + + +# --------------------------------------------------------------------------- +# Mirrors the actual train_one_step reporting math: +# +# per-rollout-mean path: +# reported = sum_of_reducer_per_mb / step_global_batch_size +# per-token-loss path: +# reported = sum_of_reducer_per_mb / sum_of_per_mb_num_tokens +# +# The reducer is the same callable used at train time (and inside +# log_rollout_data on the rollout side). +# --------------------------------------------------------------------------- + + +# 4 samples: rollout R0 owns indices 0,1,2 (mask sums 3+3+3=9); rollout R1 +# owns index 3 (mask sum 3). Pre-computed per-sample denom = group sum. +# Per-rollout-mean: R0 = 5, R1 = 11, sum = 16, divided by 2 rollouts → 8. +# Per-token-loss: sum of all x = 78, total clamped mask = 12, → 6.5. +_FIXED_RESPONSE_LENGTHS = [3, 3, 3, 3] +_FIXED_TOTAL_LENGTHS = [r + 4 for r in _FIXED_RESPONSE_LENGTHS] +_FIXED_LOSS_MASKS = [torch.ones(r, dtype=torch.float32) for r in _FIXED_RESPONSE_LENGTHS] +_FIXED_ROLLOUT_DENOMS = [9.0, 9.0, 9.0, 3.0] +_FIXED_X_PER_SAMPLE = [ + torch.tensor([1.0, 2.0, 3.0]), + torch.tensor([4.0, 5.0, 6.0]), + torch.tensor([7.0, 8.0, 9.0]), + torch.tensor([10.0, 11.0, 12.0]), +] +_FIXED_STEP_GBS = 2 # 2 distinct rollouts in the step +_EXPECTED_PER_ROLLOUT_MEAN_REPORT = 8.0 +_EXPECTED_PER_TOKEN_LOSS_REPORT = 78.0 / 12.0 + + +# Each entry: list of "rank"s, each rank is a list of mbs, each mb is the +# sample-index list packed into that mb. Covers: single mb, evenly split by +# rollout, split inside a rollout (R0 across mbs), uneven distribution, and +# fully singleton mbs per rank. +_PARTITION_CONFIGS = [ + [[[0, 1, 2, 3]]], # 1 rank, 1 mb + [[[0, 1, 2], [3]]], # 1 rank, 2 mbs split at rollout boundary + [[[0, 1], [2, 3]]], # 1 rank, 2 mbs splitting R0 across them — the tricky case + [[[0, 1]], [[2, 3]]], # 2 ranks, 1 mb each + [[[0, 1, 3]], [[2]]], # 2 ranks, R0 split across BOTH ranks (worst case for split-across-mb bug) + [[[0]], [[1]], [[2]], [[3]]], # 4 ranks, 1 sample per rank +] + + +def _simulate_report(partition, *, per_token_loss: bool) -> float: + """Reproduce train_one_step's reporting math for one partition config.""" + metric_sum = 0.0 + num_tokens_sum = 0 + for rank_mbs in partition: + for mb_indices in rank_mbs: + mb_total = [_FIXED_TOTAL_LENGTHS[i] for i in mb_indices] + mb_resp = [_FIXED_RESPONSE_LENGTHS[i] for i in mb_indices] + mb_masks = [_FIXED_LOSS_MASKS[i] for i in mb_indices] + mb_x = torch.cat([_FIXED_X_PER_SAMPLE[i] for i in mb_indices]) + if per_token_loss: + # Per-token-loss: caller uses ``calculate_per_token_loss=True`` + # to get ``sum_of_token`` (no per-sample denom). + reducer = get_sum_of_sample_mean(mb_total, mb_resp, mb_masks, calculate_per_token_loss=True) + num_tokens_sum += sum(max(int(m.sum().item()), 1) for m in mb_masks) + else: + mb_denoms = torch.tensor([_FIXED_ROLLOUT_DENOMS[i] for i in mb_indices], dtype=torch.float32) + reducer = get_sum_of_sample_mean(mb_total, mb_resp, mb_masks, mb_denoms) + metric_sum += reducer(mb_x).item() + if per_token_loss: + return metric_sum / num_tokens_sum + return metric_sum / _FIXED_STEP_GBS + + +@pytest.mark.unit +@pytest.mark.parametrize("partition", _PARTITION_CONFIGS) +def test_per_rollout_mean_report_invariant_to_mb_distribution(partition): + """Same samples should yield the same per-rollout-mean report regardless of + how they're spread across DP ranks / micro-batches — this is what lets us + change parallelism without changing wandb numbers.""" + assert _simulate_report(partition, per_token_loss=False) == pytest.approx(_EXPECTED_PER_ROLLOUT_MEAN_REPORT) + + +@pytest.mark.unit +@pytest.mark.parametrize("partition", _PARTITION_CONFIGS) +def test_per_token_loss_report_invariant_to_mb_distribution(partition): + """Same invariant for the per-token-loss reporting path.""" + assert _simulate_report(partition, per_token_loss=True) == pytest.approx(_EXPECTED_PER_TOKEN_LOSS_REPORT) + + +def _simulate_rollout_report(samples_per_rank): + """Reproduce log_rollout_data + gather_log_data's averaging math for the + per-token metric branch. + + Each "rank" applies the reducer once over its full sample subset, then + ``rollout_log_metric_contribution`` (the same helper data.py uses) emits + the ``(per_rank_sum, count)`` tuple. We aggregate via + ``Σsum / Σcount`` — the same shape ``gather_log_data`` uses. + """ + dp_size = len(samples_per_rank) + pairs: list[tuple[float, float]] = [] + for indices in samples_per_rank: + if not indices: + pairs.append( + rollout_log_metric_contribution( + 0.0, cp_size=1, num_rollouts_in_rollout=_FIXED_STEP_GBS, dp_size=dp_size + ) + ) + continue + tl = [_FIXED_TOTAL_LENGTHS[i] for i in indices] + rl = [_FIXED_RESPONSE_LENGTHS[i] for i in indices] + masks = [_FIXED_LOSS_MASKS[i] for i in indices] + denoms = torch.tensor([_FIXED_ROLLOUT_DENOMS[i] for i in indices], dtype=torch.float32) + x = torch.cat([_FIXED_X_PER_SAMPLE[i] for i in indices]) + reducer = get_sum_of_sample_mean(tl, rl, masks, denoms) + pairs.append( + rollout_log_metric_contribution( + reducer(x).item(), + cp_size=1, + num_rollouts_in_rollout=_FIXED_STEP_GBS, + dp_size=dp_size, + ) + ) + total_sum = sum(p[0] for p in pairs) + total_count = sum(p[1] for p in pairs) + return total_sum / total_count + + +_DP_PARTITIONS = [ + [[0, 1, 2, 3]], # 1 rank holds everything + [[0, 1, 2], [3]], # 2 ranks, balanced by rollout + [[0, 1], [2, 3]], # 2 ranks splitting R0 across mb-and-rank + [[0, 1, 3], [2]], # 2 ranks with R0 spread across BOTH (one of R0's samples is on rank 1) + [[0], [1], [2], [3]], # 4 ranks, one sample each (R0's samples spread across 3 ranks) +] + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_partition", _DP_PARTITIONS) +def test_rollout_report_matches_train_report_in_single_step(dp_partition): + """In a 1-step rollout, the rollout-side report (log_rollout_data → gather) + must equal the train-side report (train_one_step ``value / step_global_batch_size``) + for the same samples — otherwise wandb numbers between phases drift. + + Both go through the same reducer with the same precomputed denominators; + the contract this test pins is that the gather count plumbing on the + rollout side sums to the same denominator the train side uses + (``step_global_batch_size``), independent of how the rollout's samples + are spread across DP ranks. + """ + rollout_report = _simulate_rollout_report(dp_partition) + assert rollout_report == pytest.approx(_EXPECTED_PER_ROLLOUT_MEAN_REPORT) + + +@pytest.mark.unit +def test_train_one_step_per_rollout_mean_report_invariant_to_cp(monkeypatch, mock_dp_with_cp_group): + """End-to-end check of train_one_step's report formula across CP sizes. + + Mirrors the actual reduction order: + 1. Each (DP, CP) rank computes per-mb reducer output. + 2. Per-rank values are summed across mbs locally. + 3. All-reduce sums across DP*CP ranks. + 4. ``reduce_train_step_metrics`` applied (the same helper + ``train_one_step`` calls, so this test stays honest if the + implementation changes). + + cp_size = 1 vs cp_size = 2 must give the same reported number — + otherwise wandb metrics would drift the moment a user enables CP. + """ + from megatron.core import mpu as _mpu + + total_lengths = [12, 12] + response_lengths = [8, 8] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + sample_denoms = torch.tensor([16.0, 16.0], dtype=torch.float32) + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + step_global_batch_size = 1 # one rollout in the step + + def simulate(cp_size: int) -> float: + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: cp_size) + # values[0] is the per-token-loss path's num_tokens slot; for + # per-rollout-mean it's a zero placeholder (loss_function sets 0). + value_after_allreduce = 0.0 + for cp_rank in range(cp_size): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + if cp_size == 1: + x_for_rank = torch.cat(x_full) + else: + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + 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] + x_chunks_per_sample.append(torch.cat([c0, c1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + value_after_allreduce += reducer(x_for_rank).item() + reduced = reduce_train_step_metrics( + [{"keys": ["metric"], "values": torch.tensor([0.0, value_after_allreduce])}], + calculate_per_token_loss=False, + step_global_batch_size=step_global_batch_size, + cp_size=cp_size, + dp_with_cp_group=mock_dp_with_cp_group, + ) + return reduced["metric"] + + assert simulate(1) == pytest.approx(simulate(2)) + + +@pytest.mark.unit +def test_train_one_step_per_token_loss_report_invariant_to_cp(monkeypatch, mock_dp_with_cp_group): + """Same end-to-end check for the per-token-loss path: divisor is + ``values[0] = num_tokens`` (computed in loss.py from FULL loss masks), + which each CP rank duplicates and all-reduce sums by ``cp_size``. The + ``cp_factor = cp_size`` multiplier inside ``reduce_train_step_metrics`` + cancels that inflation, so the report stays CP-invariant. + """ + from megatron.core import mpu as _mpu + + total_lengths = [12, 12] + response_lengths = [8, 8] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + num_tokens_per_mb = sum(int(m.sum().item()) for m in loss_masks) # = 16 + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + + def simulate(cp_size: int) -> float: + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: cp_size) + value_after_allreduce = 0.0 + num_tokens_after_allreduce = 0 # each CP rank reports the same num_tokens + for cp_rank in range(cp_size): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + if cp_size == 1: + x_for_rank = torch.cat(x_full) + else: + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + 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] + x_chunks_per_sample.append(torch.cat([c0, c1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer = get_sum_of_sample_mean( + total_lengths, response_lengths, loss_masks, calculate_per_token_loss=True + ) + value_after_allreduce += reducer(x_for_rank).item() + num_tokens_after_allreduce += num_tokens_per_mb + reduced = reduce_train_step_metrics( + [ + { + "keys": ["metric"], + "values": torch.tensor([num_tokens_after_allreduce, value_after_allreduce], dtype=torch.float32), + } + ], + calculate_per_token_loss=True, + step_global_batch_size=999, # unused in per-token-loss path + cp_size=cp_size, + dp_with_cp_group=mock_dp_with_cp_group, + ) + return reduced["metric"] + + assert simulate(1) == pytest.approx(simulate(2)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_metric_report_dist.py b/tests/test_metric_report_dist.py new file mode 100644 index 000000000..7535027aa --- /dev/null +++ b/tests/test_metric_report_dist.py @@ -0,0 +1,303 @@ +"""Multi-process distributed tests for the cp_utils report helpers. + +Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` (gloo +backend) and exercise the actual production helpers end-to-end. The +single-process numerical contracts live in ``test_metric_report.py``; this +file pins the cross-rank plumbing. + +Mapping to the production train_one_step / log_rollout_data flows: + + - ``_train_step_distributed_worker`` mirrors ``train_one_step``: + per-rank reducer → ``reduce_train_step_metrics`` + (which calls ``dist.all_reduce`` over the dp-with-cp group and + applies the cp_size cancellation for the per-token-loss path). + - ``_rollout_log_distributed_worker`` mirrors ``log_rollout_data``: + per-rank reducer → ``rollout_log_metric_contribution`` → + ``gather_and_reduce_log_dict`` (which calls ``dist.gather_object`` + and applies per-key reductions). + +ALL (dp, cp) configurations must give the same reported number — that's +the contract a user touches when they flip any parallelism dial. +""" + +from __future__ import annotations + +# IMPORTANT: import the helpers (and the megatron stub it installs) BEFORE +# any slime import. Spawned workers re-import this module from scratch, so +# the same ordering must hold there — see ``stub_megatron_in_worker`` +# for the worker-side details. pytest's prepend importmode puts +# ``tests/`` on sys.path so the bare-name import works without an +# ``__init__.py``; mp.spawn children inherit the parent's sys.path. +import _cp_dist_helpers +import pytest +import torch +from _cp_dist_helpers import ( + FOUR_ROLLOUT_EXPECTED_REPORT, + FOUR_ROLLOUT_RESPONSE_LENGTHS, + FOUR_ROLLOUT_TOTAL_LENGTHS, + FOUR_ROLLOUT_X_VALUES, + cp_chunk_response_tensor, + free_port, + init_worker_process_group, + stub_megatron_in_worker, +) + + +def _train_step_distributed_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + per_token_loss: bool, + master_port: int, + result_path: str, +) -> None: + """Per-rank entrypoint for ``mp.spawn``: init gloo pg, run one rank's + share of the train-step report, write rank-0's result to a file.""" + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_with_cp_group = init_worker_process_group(rank, world_size, master_port) + try: + # Import AFTER the megatron stub override so cp_utils still binds + # against the pre-installed stub (which we've now pinned for this + # worker's CP rank). + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean, reduce_train_step_metrics + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + step_global_batch_size = 4 # 4 rollouts in the step + + # Round-robin DP partition: with 4 samples / dp=N, rank i gets + # samples i, i+N, ... (matches what _split_train_data_by_dp does + # for evenly divisible cases). + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x = [all_x[i] for i in my_indices] + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x, strict=True)] + ) + + if per_token_loss: + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, calculate_per_token_loss=True) + # num_tokens is computed off the FULL mask (not the chunked + # one) in loss.py — every CP rank reports the same number, + # which is why ``reduce_train_step_metrics`` cancels by + # ``cp_factor = cp_size`` afterwards. + num_tokens = sum(int(m.sum().item()) for m in my_masks) + values_tensor = torch.tensor([float(num_tokens), reducer(x_for_rank).item()], dtype=torch.float32) + else: + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + values_tensor = torch.tensor([0.0, reducer(x_for_rank).item()], dtype=torch.float32) + + reduced = reduce_train_step_metrics( + [{"keys": ["metric"], "values": values_tensor}], + calculate_per_token_loss=per_token_loss, + step_global_batch_size=step_global_batch_size, + cp_size=cp_size, + dp_with_cp_group=dp_with_cp_group, + ) + + if rank == 0: + with open(result_path, "w") as f: + f.write(repr(reduced["metric"])) + finally: + _dist.destroy_process_group() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_train_step_per_rollout_mean_real_distributed(dp_size, cp_size, tmp_path): + """End-to-end multi-process: spawn ``dp_size * cp_size`` workers, each + runs its share with real ``torch.distributed`` (gloo); ALL parallelism + combinations must give the same reported per-rollout-mean number. + + Expected = sum of per-rollout token-means / step_gbs + = (4.5 + 45 + 450 + 4500) / 4 = 1249.875 + """ + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.txt") + mp.spawn( + _train_step_distributed_worker, + args=(world_size, cp_size, dp_size, False, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + result = float(f.read()) + assert result == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_train_step_per_token_loss_real_distributed(dp_size, cp_size, tmp_path): + """Same end-to-end multi-process check for the per-token-loss path. + + Expected = sum of all x / total_tokens + = (36 + 360 + 3600 + 36000) / 32 = 1249.875 + """ + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.txt") + mp.spawn( + _train_step_distributed_worker, + args=(world_size, cp_size, dp_size, True, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + result = float(f.read()) + assert result == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + + +def _rollout_log_distributed_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + master_port: int, + result_path: str, +) -> None: + """Per-rank entrypoint for ``mp.spawn``: build a multi-key log_dict + covering all three reduction modes ``gather_and_reduce_log_dict`` + supports, run real ``dist.gather_object``, have rank 0 dump the + reduced dict via pickle for the parent to assert on. + """ + import pickle + + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_group = init_worker_process_group(rank, world_size, master_port) + try: + from vime.backends.megatron_utils.cp_utils import ( + gather_and_reduce_log_dict, + get_sum_of_sample_mean, + rollout_log_metric_contribution, + ) + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + num_rollouts_in_rollout = 4 + + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x = [all_x[i] for i in my_indices] + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x, strict=True)] + ) + + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + per_rank_reducer_sum = reducer(x_for_rank).item() + + # Exercise every reduction mode the production log_rollout_data emits. + log_dict = { + # per-rollout-mean: (sum, count) via rollout_log_metric_contribution. + # gather: Σsum / Σcount = sum_DP_full / num_rollouts. + "logp_per_rollout": rollout_log_metric_contribution( + per_rank_reducer_sum, + cp_size=cp_size, + num_rollouts_in_rollout=num_rollouts_in_rollout, + dp_size=dp_size, + ), + # per-sample-mean: (Σval, num_samples) — matches the + # ``total_lengths`` style in log_rollout_data. gather: Σsum/Σcount + # = total / total_samples = per-sample mean of total_lengths. + "total_lengths_per_sample": (float(sum(my_tl)), float(len(my_tl))), + # mean-across-ranks: plain scalar — matches log_multi_turn_data + # style. gather: Σvalue / dp_world. + "rank_local_mean": float(sum(my_tl)) / len(my_tl), + } + + reduced = gather_and_reduce_log_dict(log_dict, dp_size=world_size, dp_src_rank=0, dp_group=dp_group) + + if rank == 0: + with open(result_path, "wb") as f: + pickle.dump(reduced, f) + finally: + _dist.destroy_process_group() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_rollout_log_real_distributed_multi_key(dp_size, cp_size, tmp_path): + """End-to-end multi-process for ``gather_and_reduce_log_dict``. + + Covers the three key shapes ``log_rollout_data`` produces: + - per-rollout-mean ((sum, count) via ``rollout_log_metric_contribution``) + - per-sample-mean ((Σval, num_samples) tuple — e.g. ``total_lengths``) + - mean-across-ranks (plain float — e.g. multi_turn stats) + + All (dp, cp) configs must yield the same reduced numbers; matches the + expected values written in pure Python from the fixture. In particular + the per-rollout-mean number must equal what the train-step report tests + above land on (FOUR_ROLLOUT_EXPECTED_REPORT), pinning the cross-phase + contract. + """ + import pickle + + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.pkl") + mp.spawn( + _rollout_log_distributed_worker, + args=(world_size, cp_size, dp_size, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path, "rb") as f: + reduced = pickle.load(f) + + # per-rollout-mean: matches the train-side report — 1249.875. + assert reduced["logp_per_rollout"] == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + # per-sample-mean: every sample has total_length=12, so the average is 12. + assert reduced["total_lengths_per_sample"] == pytest.approx(12.0) + # mean-across-ranks: every rank's local mean is 12, so cross-rank mean is 12. + assert reduced["rank_local_mean"] == pytest.approx(12.0) + + +# Keep an explicit reference to silence "unused import" complaints while +# documenting that importing the helpers module is load-bearing (it +# installs the megatron stub before slime is touched). +_ = _cp_dist_helpers + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 165139dce..595a540ee 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -71,9 +71,6 @@ def init( self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True) dist.barrier(group=get_gloo_group()) - self.train_parallel_config = { - "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False), - } dist.barrier(group=get_gloo_group()) if args.offload_train: @@ -85,6 +82,20 @@ def init( args, role ) + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 + if vpp_size > 1: + from megatron.core.utils import get_model_config + + microbatch_group_size_per_vp_stage = get_model_config(self.model[0]).microbatch_group_size_per_vp_stage + else: + microbatch_group_size_per_vp_stage = 1 + self.train_parallel_config = { + "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False), + "cp_size": mpu.get_context_parallel_world_size(), + "vpp_size": vpp_size, + "microbatch_group_size_per_vp_stage": microbatch_group_size_per_vp_stage, + } + start_rollout_id = loaded_rollout_id + 1 if role == "critic": @@ -203,6 +214,12 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: rollout_data["loss_masks"] = [ torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"] ] + if "rollout_mask_sums" in rollout_data: + # Promote precomputed per-rollout mask totals to GPU tensors here + # (matching loss_masks) so the loss reducer can just divide. + rollout_data["rollout_mask_sums"] = torch.tensor( + rollout_data["rollout_mask_sums"], dtype=torch.float32, device=torch.cuda.current_device() + ) if "multimodal_train_inputs" in rollout_data: # Move multimodal training tensors to GPU in advance rollout_data["multimodal_train_inputs"] = [ @@ -385,7 +402,9 @@ def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): """Train critic and return CPU values (used as old-values for the next actor train).""" - data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + data_iterator = get_data_iterator(rollout_data) + num_microbatches = rollout_data["num_microbatches"] + global_batch_sizes = rollout_data["global_batch_sizes"] # Compute current critic values (used as old_values for value loss and for actor advantages). rollout_data.update(forward_only(get_values, self.args, self.model, data_iterator, num_microbatches)) @@ -400,6 +419,7 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): self.opt_param_scheduler, data_iterator, num_microbatches, + global_batch_sizes, ) if mpu.is_pipeline_last_stage() and "values" in rollout_data: @@ -410,7 +430,9 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data=None) -> None: # Create data iterator for log_probs and train. - data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + data_iterator = get_data_iterator(rollout_data) + num_microbatches = rollout_data["num_microbatches"] + global_batch_sizes = rollout_data["global_batch_sizes"] if self.args.use_rollout_routing_replay: self.fill_routing_replay(data_iterator, num_microbatches, rollout_data) @@ -507,6 +529,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data self.opt_param_scheduler, data_iterator, num_microbatches, + global_batch_sizes, ) self.prof.step(rollout_id=rollout_id) diff --git a/vime/backends/megatron_utils/cp_utils.py b/vime/backends/megatron_utils/cp_utils.py index b0eab9b6a..448c154c6 100644 --- a/vime/backends/megatron_utils/cp_utils.py +++ b/vime/backends/megatron_utils/cp_utils.py @@ -54,21 +54,37 @@ def get_sum_of_sample_mean( total_lengths: list[int], response_lengths: list[int], loss_masks: list[torch.Tensor], + sample_denoms: list[torch.Tensor] | torch.Tensor | None = None, calculate_per_token_loss: bool = False, qkv_format: str = "thd", max_seq_lens: list[int] | None = None, ) -> Callable[[torch.Tensor], torch.Tensor]: """ - Calculate correct sample mean for CP + Calculate correct sample mean for CP. + + The default (``sample_denoms=None``) is the legacy per-sample mean: each + sample's denominator is its own ``loss_mask.sum()``. Callers that want a + per-rollout token-weighted mean pass pre-computed per-sample denominators + (already as GPU tensors — see actor side) where every sample in the same + rollout group carries the same value (the sum of that rollout's mask + totals across every sibling sample in the step). Pre-computing at the + step level rather than per-mb is required — otherwise a rollout whose + samples land in different micro-batches would get a partial denominator + on each side. """ + if sample_denoms is None: + sample_denoms = [m.sum() for m in loss_masks] + cp_size = mpu.get_context_parallel_world_size() if cp_size == 1: def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( [ - (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1) - for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + (x_i * loss_mask_i).sum() / torch.clamp_min(denom, 1) + for x_i, loss_mask_i, denom in zip( + x.split(response_lengths, dim=0), loss_masks, sample_denoms, strict=False + ) ] ) @@ -100,9 +116,9 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( [ - (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) - for x_i, chunked_loss_mask, loss_mask in zip( - x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=False + (x_i * chunked_loss_mask).sum() / torch.clamp_min(denom, 1) + for x_i, chunked_loss_mask, denom in zip( + x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, sample_denoms, strict=False ) ] ) @@ -120,6 +136,114 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token +def reduce_train_step_metrics( + losses_reduced: list[dict], + *, + calculate_per_token_loss: bool, + step_global_batch_size: int, + cp_size: int, + dp_with_cp_group, +) -> dict[str, float]: + """Aggregate per-mb log dicts into the dict ``train_one_step`` reports. + + Pipeline (1:1 with what the train loop used to do inline): + 1. Sum each metric's per-mb ``values`` tensor locally on this rank. + 2. All-reduce across the DP*CP group (``dp_with_cp_group``). + 3. Apply the per-mode divisor / cp_factor: + - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``, + CP-inflated by ``cp_size`` because every CP rank computes the same + num_tokens off the FULL (not chunked) masks; the + ``cp_factor = cp_size`` multiplier cancels that inflation, leaving + the genuine per-token average. + - per-rollout-mean: divisor = constant ``step_global_batch_size`` from + the rollout side, never all-reduced, so no CP inflation to cancel + and ``cp_factor = 1``. + + Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce`` + to a no-op, then pre-aggregate virtual ranks themselves — this exercises + the same call shape as production while staying single-process. + """ + keys = losses_reduced[0]["keys"] + values = None + for x in losses_reduced: + values = x["values"] if values is None else values + x["values"] + assert len(keys) + 1 == values.numel() + dist.all_reduce(values, group=dp_with_cp_group) + values = values.tolist() + + if calculate_per_token_loss: + num_samples_or_tokens = values[0] + cp_factor = cp_size + else: + num_samples_or_tokens = step_global_batch_size + cp_factor = 1 + return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)} + + +def rollout_log_metric_contribution( + per_rank_reducer_sum: float, + *, + cp_size: int, + num_rollouts_in_rollout: int, + dp_size: int, +) -> tuple[float, float]: + """``(sum, count)`` tuple to hand the gather step for a per-rollout-mean + metric on the rollout side (``log_rollout_data``). + + Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout`` + (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size * + cp_size`` ranks, and each rank emits the same ``count``, so the totals + cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount = + sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports + for the same samples (when ``num_steps_per_rollout == 1``). + + Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end + in tests (single helper call per rank, returns the reduced number on + the source rank). + """ + sum_value = cp_size * per_rank_reducer_sum + count = num_rollouts_in_rollout / dp_size + return sum_value, count + + +def gather_and_reduce_log_dict( + log_dict: dict, + *, + dp_size: int, + dp_src_rank: int, + dp_group, +) -> dict | None: + """``dist.gather_object`` per-rank log_dicts + per-key reduction. + + Per key in the gathered dicts: + - ``(sum, count)`` tuple → ``Σsum / Σcount`` (per-rollout-mean shape; + pair with :func:`rollout_log_metric_contribution`). + - plain value → ``Σ / dp_size`` (legacy mean-across-ranks; the only + correct answer when ranks hold the same data). + + Returns the reduced dict on ``dp_src_rank``, ``None`` elsewhere. The + caller adds whatever metric-name prefix / wandb plumbing it wants — + this helper stays free of side effects so CPU multi-process unit tests + can drive it directly with real ``torch.distributed``. + """ + if dist.get_rank() == dp_src_rank: + gathered = [None] * dp_size + dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group) + reduced: dict = {} + for key in log_dict: + values = [d[key] for d in gathered] + first = values[0] + if isinstance(first, tuple) and len(first) == 2: + total_sum = sum(v[0] for v in values) + total_count = sum(v[1] for v in values) + reduced[key] = total_sum / total_count if total_count else 0.0 + else: + reduced[key] = sum(values) / dp_size + return reduced + dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group) + return None + + def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor: """ Gather tensors across all ranks in the context parallel group. diff --git a/vime/backends/megatron_utils/data.py b/vime/backends/megatron_utils/data.py index de6c62d55..1627537b4 100644 --- a/vime/backends/megatron_utils/data.py +++ b/vime/backends/megatron_utils/data.py @@ -10,14 +10,17 @@ from megatron.core.packed_seq_params import PackedSeqParams from vime.utils import train_metric_utils -from vime.utils.data import get_minimum_num_micro_batch_size from vime.utils.flops_utils import calculate_fwd_flops from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step -from vime.utils.seqlen_balancing import get_seqlen_balanced_partitions from vime.utils.types import RolloutBatch from ...utils import logging_utils -from .cp_utils import get_sum_of_sample_mean, slice_with_cp +from .cp_utils import ( + gather_and_reduce_log_dict, + get_sum_of_sample_mean, + rollout_log_metric_contribution, + slice_with_cp, +) logger = logging.getLogger(__name__) @@ -53,9 +56,6 @@ def get_batch( assert "tokens" in keys batch = data_iterator.get_next(keys) - if "dynamic_global_batch_size" in data_iterator.rollout_data: - batch["dynamic_global_batch_size"] = data_iterator.rollout_data["dynamic_global_batch_size"] - tokens = batch["tokens"] # use 0 as the pad token id should be fine? pad_token_id = 0 @@ -180,105 +180,69 @@ def gather_log_data( metric_name: str, args: Namespace, rollout_id: int, - log_dict: dict[str, float], + log_dict: dict[str, "float | tuple[float, float]"], ) -> dict[str, float] | None: """ - Gather per-rank metrics, reduce by mean on the DP source rank, and log. - - Expects `log_dict` to contain plain scalars. The DP source rank prints and - optionally logs to WandB/TensorBoard with a step derived from `rollout_id` and - batch sizes. Returns the reduced dict on the DP source rank; returns None on others. - """ - - if mpu.get_data_parallel_rank(with_context_parallel=True) == 0: - dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True) - - gathered_log_dict = [None] * dp_size - # Not sure if this will be a performance bottleneck. - dist.gather_object( - log_dict, - gathered_log_dict, - dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), - group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), - ) - - reduced_log_dict = { - f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict - } - logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") + Gather per-rank metrics, reduce on the DP source rank, and log to W&B / TB. - # Calculate step once to avoid duplication - step = compute_rollout_step(args, rollout_id) - reduced_log_dict["rollout/step"] = step - logging_utils.log(args, reduced_log_dict, step_key="rollout/step") + Each value in ``log_dict`` is either: + * a ``(sum, count)`` tuple → reduced as ``Σsum / Σcount``; + * a plain scalar → reduced as ``Σ / dp_size`` (mean across ranks). - return reduced_log_dict - else: - dist.gather_object( - log_dict, - None, - dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), - group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), - ) + The gather + reduce step is delegated to + :func:`cp_utils.gather_and_reduce_log_dict` so it can be exercised by + CPU multi-process unit tests directly. This function adds the + ``metric_name`` prefix and the W&B / TB logging side effects. + """ + reduced = gather_and_reduce_log_dict( + log_dict, + dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True), + dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True), + dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), + ) + if reduced is None: return None + reduced_log_dict = {f"{metric_name}/{k}": v for k, v in reduced.items()} + logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") + # Calculate step once to avoid duplication + step = compute_rollout_step(args, rollout_id) + reduced_log_dict["rollout/step"] = step + logging_utils.log(args, reduced_log_dict, step_key="rollout/step") + return reduced_log_dict class DataIterator: - """Micro-batch iterator over rollout dicts. - - Supports either fixed contiguous micro-batches or an explicit per-step - index schedule (for dynamic batch sizing / sequence-length balancing). - """ + """Iterator over a rollout dict following an explicit micro-batch index schedule.""" def __init__( self, rollout_data: RolloutBatch, - micro_batch_size: int | None = None, - micro_batch_indices: list[list[int]] | None = None, + micro_batch_indices: list[list[int]], ) -> None: - """Initialize an iterator over `rollout_data`. + """Initialize an iterator over ``rollout_data``. Args: - rollout_data: Dict of per-sample fields for the local step. - micro_batch_size: Fixed contiguous slice size when not using dynamic scheduling. - micro_batch_indices: Explicit indices per micro-batch when using dynamic balancing. - Must be mutually exclusive with `micro_batch_size`. + rollout_data: Dict of per-sample fields for this DP rank. + micro_batch_indices: List of mbs, each mbs being the local sample indices to select. """ self.rollout_data = rollout_data - self.micro_batch_size = micro_batch_size self.micro_batch_indices = micro_batch_indices - assert micro_batch_size is None or micro_batch_indices is None self.offset = 0 def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: """Return the next micro-batch for the requested keys. - - If `micro_batch_indices` is provided, selects rows according to the current - index list for each requested key. - - Otherwise, slices a contiguous window of size `micro_batch_size` starting - at the current offset. - Returns a dict mapping each key to a list subset (or None if absent). """ batch = {} + indices = self.micro_batch_indices[self.offset] for key in keys: vals = self.rollout_data.get(key, None) if vals is None: batch[key] = None else: - if self.micro_batch_indices is not None: - indices = self.micro_batch_indices[self.offset] - batch[key] = [vals[i] for i in indices] - else: - assert self.offset + self.micro_batch_size <= len( - vals - ), f"offset: {self.offset}, micro_batch_size: {self.micro_batch_size}, len(vals): {len(vals)}" - batch[key] = vals[self.offset : self.offset + self.micro_batch_size] - - if self.micro_batch_indices is not None: - self.offset += 1 - else: - self.offset += self.micro_batch_size + batch[key] = [vals[i] for i in indices] + self.offset += 1 return batch def reset(self) -> "DataIterator": @@ -287,102 +251,11 @@ def reset(self) -> "DataIterator": return self -def get_data_iterator( - args: Namespace, - model: torch.nn.Module | Sequence[torch.nn.Module], - rollout_data: RolloutBatch, -) -> tuple[list[DataIterator], list[int]]: - """ - Create iterators and a micro-batch schedule for a rollout step. - - - If `use_dynamic_batch_size` is False, splits into fixed-size contiguous - micro-batches of `micro_batch_size`. - - If True, computes the number of micro-batches per local step based on - `max_tokens_per_gpu` and per-sample lengths, all-reduces to a DP-wide - maximum, optionally enforces divisibility for Virtual Pipeline Parallelism (VPP), and builds a balanced - index schedule to equalize token counts across micro-batches. - - Returns `(data_iterators, num_microbatches)` where: - - `data_iterators`: list of `DataIterator`, one per VPP stage (size 1 if VPP disabled) - - `num_microbatches`: list[int], one per local step in the rollout (length = steps) - """ - dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) - dp_group = mpu.get_data_parallel_group() - vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() - if vpp_size is None: - vpp_size = 1 - if vpp_size > 1: - from megatron.core.utils import get_model_config - - config = get_model_config(model[0]) - microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage - cp_size = mpu.get_context_parallel_world_size() - - num_local_samples = len(rollout_data["total_lengths"]) - global_batch_size = rollout_data.get("dynamic_global_batch_size", args.global_batch_size) - num_local_gbs = global_batch_size // dp_size - num_steps_per_rollout = num_local_samples // num_local_gbs - - if global_batch_size != args.global_batch_size: - logger.info( - f"Using dynamic global_batch_size={global_batch_size} (original={args.global_batch_size}), " - f"num_local_samples={num_local_samples}, num_steps_per_rollout={num_steps_per_rollout}" - ) - - def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices=None): - data_iterator = [] - for _ in range(vpp_size): - data_iterator.append(DataIterator(rollout_data, micro_batch_size, micro_batch_indices)) - return data_iterator - - if not args.use_dynamic_batch_size: - num_microbatches = [num_local_gbs // args.micro_batch_size for _ in range(num_steps_per_rollout)] - data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size) - else: - assert args.max_tokens_per_gpu is not None - # calculate the number of mirobatches for each step - samples = rollout_data["total_lengths"] - assert len(samples) == num_local_samples - num_microbatches = [] - for i in range(num_steps_per_rollout): - start, end = i * num_local_gbs, (i + 1) * num_local_gbs - num_microbatches.append( - get_minimum_num_micro_batch_size(samples[start:end], args.max_tokens_per_gpu * cp_size) - ) - - num_microbatches = torch.tensor(num_microbatches, dtype=torch.int, device=torch.cuda.current_device()) - dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=dp_group) - - if vpp_size > 1: - # vpp requies the number of microbatches to be divisible by vpp_size - num_microbatches = torch.clamp( - num_microbatches // microbatch_group_size_per_vp_stage * microbatch_group_size_per_vp_stage, - min=1, - ) - - num_microbatches = num_microbatches.tolist() - - # balance the each micro batch - samples = rollout_data["total_lengths"] - # balance the number of mirobatches across steps - micro_batch_indices = [] - for i, num_mbs in enumerate(num_microbatches): - start, end = i * num_local_gbs, (i + 1) * num_local_gbs - samples = rollout_data["total_lengths"][start:end] - partitions = get_seqlen_balanced_partitions(samples, num_mbs, equal_size=False) - for j in range(num_mbs): - for k in range(len(partitions[j])): - partitions[j][k] += start - micro_batch_indices.extend(partitions) - - assert len(set(sum(micro_batch_indices, []))) == num_local_samples - - data_iterator = _generate_data_iterator(rollout_data, None, micro_batch_indices) - - return ( - data_iterator, - num_microbatches, - ) +def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]: + """Build one ``DataIterator`` per VPP stage from the pre-computed schedule in ``rollout_data``.""" + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 + micro_batch_indices = rollout_data["micro_batch_indices"] + return [DataIterator(rollout_data, micro_batch_indices) for _ in range(vpp_size)] def log_rollout_data( @@ -406,6 +279,16 @@ def log_rollout_data( loss_masks = rollout_data["loss_masks"] total_lengths = rollout_data["total_lengths"] max_seq_lens = rollout_data.get("max_seq_lens", None) + # Same per-rollout denominators the training loss uses, so reported + # log_probs / returns / advantages / etc. live in the same per-rollout + # mean space (rather than per-sample) as the gradient signal. + rollout_mask_sums = rollout_data.get("rollout_mask_sums", None) + # For per-rollout-mean metrics: ``rollout_log_metric_contribution`` + # produces the ``(sum, count)`` tuple so gather_log_data's + # ``Σsum / Σcount`` lands on ``sum_DP_full / num_rollouts`` — the + # same number train_one_step reports for the same samples. + dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False) + num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"]) for key, val in rollout_data.items(): if key in [ @@ -413,15 +296,20 @@ def log_rollout_data( "multimodal_train_inputs", "loss_masks", "sample_indices", + "rollout_ids", + "rollout_mask_sums", "rollout_routed_experts", "max_seq_lens", - "dynamic_global_batch_size", + "global_batch_sizes", + "num_microbatches", + "micro_batch_indices", ]: continue - # Upload per sample mean for each rollout value - # There are the following assumptions: - # - Each dp rank has the same number of samples + # Emit (sum, count) so gather_log_data can do a weighted average across + # DP ranks. This stops the legacy "every rank has the same N samples" + # assumption from biasing means once uneven-DP partitioning lands. if isinstance(val, (list, tuple)): + count = len(val) if isinstance(val[0], torch.Tensor): # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be # modified in place and will cause problem for the next rollout. @@ -435,25 +323,38 @@ def log_rollout_data( "teacher_log_probs", "opd_reverse_kl", ]: - val = torch.cat(val).clone().detach() + tensor = torch.cat(val).clone().detach() sum_of_sample_mean = get_sum_of_sample_mean( total_lengths, response_lengths, loss_masks, + rollout_mask_sums, qkv_format=args.qkv_format, max_seq_lens=max_seq_lens, ) - val = cp_size * sum_of_sample_mean(val) / len(loss_masks) - else: - val = torch.cat(val).clone().detach() - val = val.mean() * cp_size + # Compute (sum, count) via the shared helper so this + # path and the unit tests stay in sync. + sum_value, count = rollout_log_metric_contribution( + sum_of_sample_mean(tensor).item(), + cp_size=cp_size, + num_rollouts_in_rollout=num_rollouts_in_rollout, + dp_size=dp_world, + ) + log_dict[key] = (sum_value, count) + continue + tensor = torch.cat(val).clone().detach() + # val.mean() * cp_size is the per-sample mean for one rank; + # multiply by count to get the per-rank sum. + per_rank_sum = tensor.mean() * cp_size * count + sum_value = per_rank_sum.item() else: - val = sum(val) / len(val) + sum_value = sum(val) + log_dict[key] = (sum_value, count) elif isinstance(val, torch.Tensor): - val = val.float().mean() + # Scalar tensor (one per rank): treat as count=1. + log_dict[key] = (val.float().mean().item(), 1) else: raise ValueError(f"Unsupported type: {type(val)} for key: {key}") - log_dict[key] = val.item() if isinstance(val, torch.Tensor) else val reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict) if args.ci_test and reduced_log_dict is not None: @@ -531,8 +432,14 @@ def quantile(total_value, n_quantiles, data) -> dict: for p, val in correct_response_length_percentile.items(): rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses if len(correct_entropy) > 0: + # NOTE: per-sample-mean over the correct subset, not per-rollout. + # A rollout's siblings may not all be correct, and slicing + # ``rollout_mask_sums`` here would leave a denom that still + # includes incorrect siblings — meaningless for a "correct-only" + # entropy report. Per-sample-mean over the filtered subset is + # the cleanest semantic. sum_of_sample_mean = get_sum_of_sample_mean( - correct_total_lengths, correct_response_lengths, correct_loss_masks + correct_total_lengths, correct_response_lengths, correct_loss_masks, sample_denoms=None ) correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0)) rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 6f95550b2..3f6ab29c5 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -931,12 +931,17 @@ def policy_loss_function( tis_func = vanilla_tis_function pg_loss, modified_response_masks, tis_metrics = tis_func(**tis_kwargs) - # [decouple IS and rejection] Rebuild sum_of_sample_mean with modified_response_masks for denominator correction - # modified_response_masks will be sliced with cp in get_sum_of_sample_mean + # [decouple IS and rejection] Rebuild sum_of_sample_mean with + # modified_response_masks for numerator correction (rejected tokens + # zeroed in pg_loss). Denominators stay the precomputed per-rollout + # totals from ``rollout_mask_sums`` (based on original loss_masks) — + # same normalizer as the outer reducer, so pg_loss and the rest of the + # reported metrics live in the same per-rollout-mean space. sum_of_sample_mean = get_sum_of_sample_mean( total_lengths, response_lengths, modified_response_masks, + batch["rollout_mask_sums"], args.calculate_per_token_loss, args.qkv_format, max_seq_lens, @@ -1136,6 +1141,7 @@ def loss_function( args: Namespace, batch: RolloutBatch, num_microbatches: int, + step_global_batch_size: int, logits: torch.Tensor, ) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]: """Dispatch to the configured loss and rescale for Megatron integration. @@ -1147,10 +1153,14 @@ def loss_function( Args: args: Configuration specifying `loss_type`, `calculate_per_token_loss`, - `global_batch_size`, and optionally `custom_loss_function_path`. + and optionally `custom_loss_function_path`. batch: Mini-batch with "loss_masks", "response_lengths", and other keys required by the selected loss function. num_microbatches: Number of gradient accumulation steps. + step_global_batch_size: Sample count for the current training step + (total across DP). Replaces the legacy ``args.global_batch_size`` + fallback so the train side stops depending on "every DP rank holds + the same N samples". logits: Model outputs (policy or value head). Returns: @@ -1162,12 +1172,12 @@ def loss_function( "values" (1D tensor: [count, metric1, metric2, ...]). """ num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]]) - num_samples = len(batch["response_lengths"]) sum_of_sample_mean = get_sum_of_sample_mean( batch["total_lengths"], batch["response_lengths"], batch["loss_masks"], + batch["rollout_mask_sums"], args.calculate_per_token_loss, args.qkv_format, batch.get("max_seq_lens", None), @@ -1199,10 +1209,12 @@ def loss_function( loss = loss + 0 * logits.sum() # Here we need to divide by cp_size because to cancel the multiply in Megatron. - global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size) if not args.calculate_per_token_loss: loss = ( - loss * num_microbatches / global_batch_size * mpu.get_data_parallel_world_size(with_context_parallel=True) + loss + * num_microbatches + / step_global_batch_size + * mpu.get_data_parallel_world_size(with_context_parallel=True) ) else: loss = loss * mpu.get_context_parallel_world_size() @@ -1212,9 +1224,16 @@ def loss_function( (num_tokens if args.calculate_per_token_loss else torch.tensor(1, device=logits.device)), { "keys": list(log.keys()), + # values[0] is the consumer's reporting denominator after + # all-reduce. For per-token-loss it must equal step total tokens + # (only known by summing per-mb num_tokens across mbs / DP). For + # per-rollout-mean it is a constant — ``step_global_batch_size`` — + # so we leave a 0 placeholder here and let ``train_one_step`` + # substitute the constant directly, instead of routing it through + # per-mb fractions. "values": torch.tensor( [ - num_samples if not args.calculate_per_token_loss else num_tokens, + num_tokens if args.calculate_per_token_loss else 0, ] + list(log.values()), device=logits.device, diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index 1e9540c84..b650ac4b0 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -27,6 +27,7 @@ from vime.utils.memory_utils import clear_memory from .checkpoint import load_checkpoint, save_checkpoint +from .cp_utils import reduce_train_step_metrics from .data import DataIterator, get_batch from .loss import loss_function from .model_provider import get_model_provider_func @@ -143,7 +144,15 @@ def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) Returns: OptimizerParamScheduler: Initialized scheduler bound to ``optimizer``. """ - # Iteration-based training. + # Iteration-based training. ``train_iters`` is an estimate of the total + # number of training steps — it's only used to size Megatron's LR decay + # schedule (and ``lr_decay_iters`` defaults to it). With variable per-rollout + # sample counts (dynamic sampling / filtering / custom step splitter) the + # *actual* total can drift; the schedule still tracks the true progress via + # ``opt_param_scheduler.num_steps`` (samples consumed, also persisted across + # resume), so the worst case is the cosine/linear schedule reaches its + # plateau slightly early or late. Pass ``--lr-decay-iters`` explicitly if you + # need exact decay control. args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size if args.lr_decay_iters is None: args.lr_decay_iters = args.train_iters @@ -412,6 +421,7 @@ def train_one_step( optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler, num_microbatches: int, + step_global_batch_size: int, microbatch_pbar=None, ) -> tuple[dict[str, float], float]: """Execute a single pipeline-parallel training step. @@ -428,6 +438,13 @@ def train_one_step( optimizer (MegatronOptimizer): Optimizer instance. opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. num_microbatches (int): Number of microbatches to process. + step_global_batch_size (int): Rollout count for this training step + (total across DP; one "rollout" = one execution of one of the + ``n_samples_per_prompt`` rollouts, which may emit >1 training + sample under compact / subagent). Used both as the loss + normalizer inside the closure and as the LR scheduler + ``increment``. In the common case (1 rollout = 1 sample) this + equals the per-step sample count, so behavior is unchanged. Returns: tuple[dict[str, float], float]: Reduced loss dictionary (last stage only) @@ -480,6 +497,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "rollout_log_probs", "max_seq_lens", "teacher_log_probs", + "rollout_mask_sums", ], args.data_pad_size_multiplier, args.qkv_format, @@ -522,7 +540,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": os.environ["ROUTING_REPLAY_STAGE"] = old_stage - return output_tensor, partial(loss_function, args, batch, num_microbatches) + return output_tensor, partial(loss_function, args, batch, num_microbatches, step_global_batch_size) # Forward pass. forward_backward_func = get_forward_backward_func() @@ -561,9 +579,10 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p # Update parameters. update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - # Update learning rate. + # Update learning rate. Use the per-step global_batch_size when dynamic + # batching is on so the scheduler's samples-seen counter tracks reality. assert update_successful - opt_param_scheduler.step(increment=args.global_batch_size) + opt_param_scheduler.step(increment=step_global_batch_size) # release grad for model_chunk in model: @@ -571,22 +590,13 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p optimizer.zero_grad() if mpu.is_pipeline_last_stage(ignore_virtual=True): - # Average loss across microbatches. - keys = losses_reduced[0]["keys"] - values = None - for x in losses_reduced: - if values is None: - values = x["values"] - else: - values += x["values"] - assert len(keys) + 1 == values.numel() - torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True)) - - loss_reduced = {} - values = values.tolist() - num_samples_or_tokens = values[0] - for key, value in zip(keys, values[1:], strict=False): - loss_reduced[key] = value * mpu.get_context_parallel_world_size() / num_samples_or_tokens + loss_reduced = reduce_train_step_metrics( + losses_reduced, + calculate_per_token_loss=args.calculate_per_token_loss, + step_global_batch_size=step_global_batch_size, + cp_size=mpu.get_context_parallel_world_size(), + dp_with_cp_group=mpu.get_data_parallel_group(with_context_parallel=True), + ) return loss_reduced, grad_norm return {}, grad_norm @@ -603,6 +613,7 @@ def train( opt_param_scheduler: OptimizerParamScheduler, data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], + global_batch_sizes: Sequence[int], ) -> None: """Run training over a rollout consisting of multiple steps. @@ -616,9 +627,20 @@ def train( opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches. num_microbatches (Sequence[int]): Microbatches per step in the rollout. + global_batch_sizes (Sequence[int]): Rollout count per step (total + across DP; one "rollout" = one execution of one of the + ``n_samples_per_prompt`` rollouts of a prompt). Same length as + ``num_microbatches``; consumed by ``train_one_step`` for loss + scaling and LR scheduler increments. Equals per-step sample count + in the common case (1 rollout = 1 sample). """ args = get_args() + assert len(num_microbatches) == len(global_batch_sizes), ( + f"num_microbatches and global_batch_sizes must have the same length, " + f"got {len(num_microbatches)} vs {len(global_batch_sizes)}" + ) + for iterator in data_iterator: iterator.reset() @@ -713,6 +735,7 @@ def train( optimizer, opt_param_scheduler, num_microbatches[step_id], + global_batch_sizes[step_id], microbatch_pbar=microbatch_pbar, ) @@ -766,6 +789,8 @@ def train( for param_group_id, param_group in enumerate(optimizer.param_groups): log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group) + # Per-step gbs — uneven step sizes are easy to miss without this. + log_dict[f"train/{role_tag}global_batch_size"] = global_batch_sizes[step_id] log_dict["train/step"] = accumulated_step_id logging_utils.log(args, log_dict, step_key="train/step") @@ -815,7 +840,10 @@ def train( def save( - iteration: int, model: Sequence[DDP], optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler + iteration: int, + model: Sequence[DDP], + optimizer: MegatronOptimizer, + opt_param_scheduler: OptimizerParamScheduler, ) -> None: """Persist a training checkpoint safely with forward hooks disabled. diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 73db78e74..298270d62 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -21,12 +21,12 @@ GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" from vime.rollout.base_types import call_rollout_fn from vime.utils import logging_utils +from vime.utils.dp_schedule import build_dp_schedule from vime.utils.health_monitor import RolloutHealthMonitor from vime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client from vime.utils.logging_utils import configure_logger, init_tracking from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix from vime.utils.misc import Box, group_by, load_function -from vime.utils.seqlen_balancing import get_seqlen_balanced_partitions from vime.utils.types import Sample from ..utils.metric_utils import has_repetition @@ -486,7 +486,7 @@ def generate(self, rollout_id): # if debug rollout only, we don't convert samples to train data and directly return return data = self._convert_samples_to_train_data(data) - return self._split_train_data_by_dp(data, self.train_parallel_config["dp_size"]) + return self._split_train_data_by_dp(data) def eval(self, rollout_id): if self.args.debug_train_only: @@ -581,56 +581,19 @@ def _get_rollout_data(self, rollout_id): data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False) metrics = data.metrics data = data.samples + # Enforce the rollout_id contract before flattening: any list[Sample] + # encountered in the nested output must have rollout_id set on every + # element. Default rollouts inherit it from the data source; compact / + # subagent paths that split one rollout into N training samples must + # set the same rollout_id on every sibling so the loss reducer counts + # the rollout once instead of N times. + _validate_rollout_id_annotated(data) # flatten the data if it is a list of lists while isinstance(data[0], list): data = list(itertools.chain.from_iterable(data)) - if not self.args.disable_rollout_trim_samples and not self.args.debug_rollout_only: - global_batch_size = self.args.global_batch_size - if self.args.use_dynamic_global_batch_size: - logger.info(f"Collected {len(data)} samples from rollout to train with dynamic global batch size") - # TODO: this is a temporary solution, we should directly save dynamic_global_batch_size to rollout data - self._dynamic_global_batch_size = self._compute_dynamic_global_batch_size(len(data)) - global_batch_size = self._dynamic_global_batch_size - - if len(data) % global_batch_size != 0: - trim_len = (len(data) // global_batch_size) * global_batch_size - if trim_len == 0: - raise ValueError(f"Not enough samples {len(data)} for global_batch_size {global_batch_size}") - origin_data_length = len(data) - data = data[:trim_len] - logger.info(f"trim number of samples from {origin_data_length} to {trim_len}") - logger.info(f"Final collected {len(data)} samples from rollout to train") - return data, metrics - def _compute_dynamic_global_batch_size(self, num_samples: int) -> int: - """Calculate dynamic global_batch_size to ensure only one training step. - - Strategy: global_batch_size = num_samples rounded down to a multiple of dp_size - This ensures num_steps_per_rollout = num_samples // global_batch_size = 1 - """ - dp_size = self.train_parallel_config["dp_size"] - original_gbs = self.args.global_batch_size - - # Round down to a multiple of dp_size to ensure only one training step - dynamic_gbs = (num_samples // dp_size) * dp_size - - if dynamic_gbs == 0: - # Too few samples, use at least dp_size - dynamic_gbs = dp_size - logger.warning(f"num_samples={num_samples} < dp_size={dp_size}, using dp_size as global_batch_size") - - # Calculate how many samples will be discarded - wasted = num_samples - dynamic_gbs - - if dynamic_gbs != original_gbs or wasted > 0: - logger.info( - f"Dynamic global_batch_size: {original_gbs} -> {dynamic_gbs} (num_samples={num_samples}, dp_size={dp_size}, num_steps=1, wasted={wasted})" - ) - - return dynamic_gbs - def _save_debug_rollout_data(self, data, rollout_id, evaluation: bool): # TODO to be refactored (originally Buffer._set_data) if (path_template := self.args.save_debug_rollout_data) is not None: @@ -698,6 +661,12 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "raw_reward": raw_rewards, "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples], "sample_indices": [sample.index for sample in samples], + # Rollout id (one per rollout execution). Default rollouts emit one + # sample per rollout, so we fall back to ``sample.index`` (unique). + # Compact / subagent paths that emit multiple training samples per + # rollout set ``rollout_id`` explicitly so all siblings share a + # value; the loss reducer then aggregates them as one rollout. + "rollout_ids": [s.rollout_id if s.rollout_id is not None else s.index for s in samples], } # loss mask @@ -716,6 +685,23 @@ 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 + # Per-rollout aggregate, precomputed at the step level (where we can + # see every sample of every rollout) and broadcast per-sample so the + # per-mb loss reducer uses the correct whole-rollout denominator even + # when a rollout's samples land in different micro-batches (first-fit + # packing can split a rollout across mbs): + # + # ``rollout_mask_sums[i]`` — sum of loss-mask totals over every + # sample in sample i's rollout. Used as the reducer's denominator + # so summing partial contributions across mbs yields one + # token-weighted mean per rollout. + rollout_id_list = train_data["rollout_ids"] + mask_sums_per_sample = [sum(m) for m in loss_masks] + rollout_total_mask: dict[int, int] = {} + for rid, ms in zip(rollout_id_list, mask_sums_per_sample, strict=True): + rollout_total_mask[rid] = rollout_total_mask.get(rid, 0) + ms + train_data["rollout_mask_sums"] = [rollout_total_mask[rid] for rid in rollout_id_list] + # 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): @@ -749,27 +735,35 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl def set_train_parallel_config(self, config: dict): self.train_parallel_config = config - def _split_train_data_by_dp(self, data, dp_size): - """Split the train data by data parallel size.""" - rollout_data = {} - - if "prompt" in data: - rollout_data["prompt"] = data["prompt"] - + def _split_train_data_by_dp(self, data): + """Compute the DP/mbs schedule and package each rank's rollout_data + into a Ray Box. The schedule itself is computed by + :func:`build_dp_schedule` so it stays unit-testable without Ray/sglang. + + Step split is by rollout id (``samples[i].rollout_id``, falling back + to ``samples[i].index``); each step holds exactly + ``args.global_batch_size`` rollouts so the training-step count per + rollout is fixed at ``rollout_batch_size * n_samples_per_prompt // + global_batch_size`` regardless of how many training samples each + rollout produced. + """ + dp_size = self.train_parallel_config["dp_size"] total_lengths = [len(t) for t in data["tokens"]] data["total_lengths"] = total_lengths - if self.args.balance_data: - partitions = get_seqlen_balanced_partitions(total_lengths, dp_size, equal_size=True) - else: - partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] + partitions, micro_batch_indices, num_microbatches, global_batch_sizes = build_dp_schedule( + self.args, + self.train_parallel_config, + total_lengths, + global_batch_size=self.args.global_batch_size, + rollout_indices=data["rollout_ids"], + ) + # Package per-rank rollout_data rollout_data_refs = [] - - for i in range(dp_size): - rollout_data = {} - partition = partitions[i] - rollout_data["partition"] = partition + for r in range(dp_size): + partition = partitions[r] + rollout_data = {"partition": partition} for key in [ "tokens", "multimodal_train_inputs", @@ -779,6 +773,8 @@ def _split_train_data_by_dp(self, data, dp_size): "loss_masks", "round_number", "sample_indices", + "rollout_ids", + "rollout_mask_sums", "rollout_log_probs", "rollout_routed_experts", "prompt", @@ -786,23 +782,51 @@ def _split_train_data_by_dp(self, data, dp_size): ]: if key not in data: continue - val = [data[key][j] for j in partition] - rollout_data[key] = val + rollout_data[key] = [data[key][j] for j in partition] # keys that need to be splited at train side - for key in [ - "raw_reward", - "total_lengths", - ]: + for key in ["raw_reward", "total_lengths"]: if key not in data: continue rollout_data[key] = data[key] - # Pass dynamic global_batch_size to training side - if hasattr(self, "_dynamic_global_batch_size"): - rollout_data["dynamic_global_batch_size"] = self._dynamic_global_batch_size + rollout_data["global_batch_sizes"] = global_batch_sizes + rollout_data["num_microbatches"] = num_microbatches + rollout_data["micro_batch_indices"] = micro_batch_indices[r] rollout_data_refs.append(Box(ray.put(rollout_data))) return rollout_data_refs +def _validate_rollout_id_annotated(node, depth=0): + """Walk the rollout function's nested output and validate ``rollout_id`` only + when a compact / subagent pattern is detected. + + "Compact" = the rollout function wraps multiple training samples from one + rollout execution into a ``list[Sample]``. In slime's convention the + default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout) + so its leaf ``list[Sample]`` lands at depth 1 and we skip validation, + preserving backward compatibility. A compact rollout adds a third level: + ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout), + so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require + every sibling to carry a non-None ``rollout_id`` and to share the same + value, so the loss reducer counts the rollout once instead of N times. + """ + if isinstance(node, Sample): + return + assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}" + if node and isinstance(node[0], Sample): + if depth >= 2 and len(node) > 1: + rids = [s.rollout_id for s in node] + missing = [i for i, r in enumerate(rids) if r is None] + assert not missing, ( + f"Compact rollout returned {len(node)} samples but rollout_id is unset on " + f"positions {missing}. Set Sample.rollout_id on every sibling so the loss " + "reducer can aggregate them as one rollout instead of N." + ) + assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}." + return + for item in node: + _validate_rollout_id_annotated(item, depth + 1) + + def _allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): addr_and_ports = {} for rank, _ in rollout_engines: diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 05c308234..a1e4d8204 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -1305,18 +1305,6 @@ def add_rollout_buffer_arguments(parser): "can process all samples including filtered ones." ), ) - parser.add_argument( - "--disable-rollout-trim-samples", - action="store_true", - default=False, - help="disable trim samples in rollout buffer when converting samples to train data", - ) - parser.add_argument( - "--use-dynamic-global-batch-size", - action="store_true", - default=False, - help="enable dynamic global batch size, disable trim samples in rollout buffer when converting samples to train data", - ) return parser def add_custom_megatron_plugins_arguments(parser): diff --git a/vime/utils/data.py b/vime/utils/data.py index 3f1417437..d158ea627 100644 --- a/vime/utils/data.py +++ b/vime/utils/data.py @@ -282,20 +282,6 @@ def __len__(self): return len(self.samples) -def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): - # use first fit to get the number of micro batches - batches = [] - for length in total_lengths: - for i in range(len(batches)): - if batches[i] + length <= max_tokens_per_gpu: - batches[i] += length - break - else: - batches.append(length) - - return len(batches) - - def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): assert len(rollout_data_ref) == dp_size rollout_data = ray.get(rollout_data_ref[dp_rank].inner) diff --git a/vime/utils/dp_schedule.py b/vime/utils/dp_schedule.py new file mode 100644 index 000000000..e30621ca5 --- /dev/null +++ b/vime/utils/dp_schedule.py @@ -0,0 +1,191 @@ +"""Per-rollout DP/microbatch scheduling. + +Pure-Python logic that decides, for one rollout's worth of sample lengths, +how to group samples into micro-batches and which DP rank owns each mbs. +Lives outside the ray/sglang-importing modules so it can be unit-tested +under CPU-only CI. + +The scheduling philosophy is **pack first, distribute second**: + + 1. Group samples by rollout id (``rollout_indices[i]`` = + ``samples[i].index``) and split rollouts into steps of + ``global_batch_size`` rollouts each. In the common case one rollout + emits one training sample so this is the same as a contiguous chunk; + under compact / subagent one rollout may emit multiple training + samples, in which case all of those samples stay in the same step. + 2. For each step, pack its samples into ``K`` micro-batches with a + single first-fit pass (dynamic batch) or fixed-size chunking + (static batch). + 3. Adjust ``K`` to a multiple of ``dp_size * (mb_group if vpp>1 else 1)`` + by splitting the largest multi-sample bins (dynamic only). + 4. Distribute the ``K`` mbs across ``dp_size`` ranks, ``K / dp_size`` + each, with either a strided round-robin or a Karmarkar-Karp pass on + mbs token sums. + +Invariants guaranteed by :func:`build_dp_schedule` (asserted by the tests): + - every DP rank runs the **same** ``num_microbatches`` per training step + (required for PP sync); + - every mbs (dynamic path) holds ``<= max_tokens_per_gpu * cp_size`` + tokens, with one exception — an individual sample larger than that cap + lands alone in its own mbs (and that mbs is the only one allowed to + exceed the cap); + - the union of per-rank sample indices equals the set of samples kept + after trimming trailing rollouts (every kept sample placed exactly + once); + - flattening ``micro_batch_indices`` for a rank yields + ``range(num_samples_rank)`` (each rank's samples are tiled exactly + once by its mbs schedule). +""" + +from __future__ import annotations + +import logging +from typing import Any + +from vime.utils.seqlen_balancing import expand_bins_by_splitting, first_fit_pack, get_seqlen_balanced_partitions + +logger = logging.getLogger(__name__) + + +def _pack_step_into_mbs( + step_lengths: list[int], + *, + use_dynamic_batch_size: bool, + max_per_bin: int | None, + micro_batch_size: int | None, +) -> list[list[int]]: + """Group a step's samples into mbs. Returns ``mbs[k]`` = local indices into ``step_lengths``.""" + if use_dynamic_batch_size: + assert max_per_bin is not None + return first_fit_pack(step_lengths, max_per_bin) + assert micro_batch_size is not None + n = len(step_lengths) + return [list(range(i, min(i + micro_batch_size, n))) for i in range(0, n, micro_batch_size)] + + +def build_dp_schedule( + args: Any, + train_parallel_config: dict, + total_lengths: list[int], + *, + global_batch_size: int, + rollout_indices: list[int], +) -> tuple[list[list[int]], list[list[list[int]]], list[int], list[int]]: + """Compute the per-rank DP partition and micro-batch schedule. + + See module docstring for the pack-first-distribute-second strategy. + + Args: + args: Namespace with ``micro_batch_size``, ``use_dynamic_batch_size``, + ``max_tokens_per_gpu``, ``balance_data``. + train_parallel_config: ``{"dp_size", "cp_size", "vpp_size", + "microbatch_group_size_per_vp_stage"}``. + total_lengths: token count per sample, indexed globally. + global_batch_size: number of rollouts (NOT training samples) per + training step. Number of training steps = + ``num_rollouts // global_batch_size``; trailing rollouts whose + samples don't fit are dropped. + rollout_indices: rollout id for each sample (``samples[i].index``). + Samples sharing the same id are kept together in one step. + + Returns: + ``(partitions, micro_batch_indices, num_microbatches, global_batch_sizes)``. + ``global_batch_sizes[s]`` = rollout count for step s (constant + ``global_batch_size`` for every step). + """ + dp_size = train_parallel_config["dp_size"] + cp_size = train_parallel_config["cp_size"] + vpp_size = train_parallel_config["vpp_size"] + mb_group = train_parallel_config["microbatch_group_size_per_vp_stage"] + + max_per_bin = None + if args.use_dynamic_batch_size: + assert args.max_tokens_per_gpu is not None + max_per_bin = args.max_tokens_per_gpu * cp_size + + # mbs count per step must be divisible by (dp_size * mb_group_for_vpp) so + # every rank ends up with the same num_mbs and (for VPP) the per-rank mbs + # count is a multiple of mb_group. + align_to = dp_size * (mb_group if vpp_size > 1 else 1) + + # Group samples by rollout id (preserve first-occurrence order). All + # samples from one rollout stay in a single step so the per-rollout loss + # reducer is well-defined. + rollout_id_to_samples: dict[int, list[int]] = {} + for sample_pos, rid in enumerate(rollout_indices): + rollout_id_to_samples.setdefault(rid, []).append(sample_pos) + rollout_ids = list(rollout_id_to_samples.keys()) + + num_steps = len(rollout_ids) // global_batch_size + assert num_steps >= 1, ( + f"num_rollouts ({len(rollout_ids)}) < global_batch_size ({global_batch_size}); " + f"need at least one rollout per step." + ) + + partitions: list[list[int]] = [[] for _ in range(dp_size)] + micro_batch_indices: list[list[list[int]]] = [[] for _ in range(dp_size)] + num_microbatches: list[int] = [] + global_batch_sizes: list[int] = [] + + for step_i in range(num_steps): + step_rollouts = rollout_ids[step_i * global_batch_size : (step_i + 1) * global_batch_size] + sample_indices = [pos for rid in step_rollouts for pos in rollout_id_to_samples[rid]] + step_lengths = [total_lengths[i] for i in sample_indices] + global_batch_sizes.append(global_batch_size) + assert len(sample_indices) >= dp_size, ( + f"step {step_i}: {len(sample_indices)} samples < dp_size {dp_size}; " + f"each step needs at least one sample per rank." + ) + + # 1. Pack samples in this step into mbs with one global pass. + # ``step_mbs`` indices are LOCAL into ``sample_indices``. + step_mbs = _pack_step_into_mbs( + step_lengths, + use_dynamic_batch_size=args.use_dynamic_batch_size, + max_per_bin=max_per_bin, + micro_batch_size=getattr(args, "micro_batch_size", None), + ) + + # 2. Align mbs count to a multiple of ``align_to``. + target_K = max(((len(step_mbs) + align_to - 1) // align_to) * align_to, align_to) + if target_K != len(step_mbs): + if args.use_dynamic_batch_size: + expand_bins_by_splitting(step_mbs, target_K, step_lengths) + assert len(step_mbs) == target_K, ( + f"dynamic path: could only produce {len(step_mbs)} mbs after maximal splitting; " + f"need {target_K}. step {step_i} has {len(sample_indices)} samples, below the " + f"alignment threshold ({align_to})." + ) + else: + raise AssertionError( + f"static path: num_mbs ({len(step_mbs)}) is not a multiple of " + f"dp_size * mb_group ({align_to}); got " + f"step_size={len(sample_indices)}, micro_batch_size={args.micro_batch_size}, " + f"dp_size={dp_size}, mb_group={mb_group if vpp_size > 1 else 1}. " + f"Splitting static mbs would break the fixed-size invariant; adjust the config " + f"so step_size % (dp_size * micro_batch_size * mb_group) == 0." + ) + + K = len(step_mbs) + num_mbs_per_rank = K // dp_size + num_microbatches.append(num_mbs_per_rank) + + # 3. Distribute mbs across ranks: KK on mbs token sums when balance_data is on, + # otherwise a strided round-robin. Both produce ``num_mbs_per_rank`` mbs per + # rank (equal_size=True is what KK needs for PP to stay synced). + if args.balance_data: + mbs_token_sums = [sum(step_lengths[i] for i in bin_) for bin_ in step_mbs] + rank_mbs_idx = get_seqlen_balanced_partitions(mbs_token_sums, dp_size, equal_size=True) + else: + rank_mbs_idx = [list(range(r, K, dp_size)) for r in range(dp_size)] + + # 4. Build per-rank partitions (global sample indices) and micro_batch_indices + # (local indices into partitions[r]). + for r in range(dp_size): + for mbs_idx in rank_mbs_idx[r]: + mbs_locals = step_mbs[mbs_idx] # local indices into sample_indices + local_start = len(partitions[r]) + partitions[r].extend(sample_indices[i] for i in mbs_locals) + micro_batch_indices[r].append(list(range(local_start, local_start + len(mbs_locals)))) + + return partitions, micro_batch_indices, num_microbatches, global_batch_sizes diff --git a/vime/utils/seqlen_balancing.py b/vime/utils/seqlen_balancing.py index a5dd71f94..5736d8850 100644 --- a/vime/utils/seqlen_balancing.py +++ b/vime/utils/seqlen_balancing.py @@ -177,6 +177,58 @@ def _check_and_sort_partitions(partitions): return _check_and_sort_partitions(partitions) +def first_fit_pack(total_lengths, max_tokens_per_bin): + """First-fit bin packing. + + Returns ``list[list[int]]`` — each bin is a list of indices into ``total_lengths``. + Bin sums are ``<= max_tokens_per_bin`` whenever every individual ``length`` fits; + an oversized sample lands alone in its own bin with sum equal to its length. + """ + bins: list[list[int]] = [] + bin_sums: list[int] = [] + for idx, length in enumerate(total_lengths): + for j in range(len(bins)): + if bin_sums[j] + length <= max_tokens_per_bin: + bins[j].append(idx) + bin_sums[j] += length + break + else: + bins.append([idx]) + bin_sums.append(length) + return bins + + +def _split_bin_by_tokens(bin_indices: list[int], lengths) -> list[list[int]]: + """Split a bin's indices into two halves balanced by total tokens (LPT heuristic). + + Returns ``[left, right]`` where both lists together cover ``bin_indices``. Because + each half is a strict subset of ``bin_indices``, both have token sums ``<=`` the + original bin's sum — useful when you need to grow a bin packing without ever + creating a bin larger than the originals. + """ + halves: list[list[int]] = [[], []] + sums = [0, 0] + for idx in sorted(bin_indices, key=lambda i: -lengths[i]): + h = 0 if sums[0] <= sums[1] else 1 + halves[h].append(idx) + sums[h] += lengths[idx] + return halves + + +def expand_bins_by_splitting(bins: list[list[int]], target_count: int, lengths) -> None: + """Grow ``bins`` in place to ``target_count`` by repeatedly splitting the largest + multi-sample bin via :func:`_split_bin_by_tokens`. Stops early if every remaining + bin is a singleton (no bin can be split further).""" + while len(bins) < target_count: + candidates = [(sum(lengths[i] for i in b), idx) for idx, b in enumerate(bins) if len(b) > 1] + if not candidates: + break + _, idx = max(candidates) + left, right = _split_bin_by_tokens(bins[idx], lengths) + bins[idx] = left + bins.append(right) + + def get_reverse_idx(idx_map): reverse_idx_map = copy.deepcopy(idx_map) diff --git a/vime/utils/types.py b/vime/utils/types.py index 9ac3916c2..a05a2d3bd 100644 --- a/vime/utils/types.py +++ b/vime/utils/types.py @@ -11,6 +11,14 @@ class Sample: group_index: int | None = None index: int | None = None + # Id of the rollout this sample came from. Defaults to ``None`` and the + # downstream pipeline falls back to ``index`` (so the default rollout + # path, where one execution = one training sample, sees rollout_id == + # index). Compact / subagent paths that split one rollout execution into + # multiple training samples should set the same ``rollout_id`` on every + # sibling, so loss aggregation averages within the rollout instead of + # over-counting it. + rollout_id: int | None = None # prompt prompt: str | list[dict[str, str]] = "" tokens: list[int] = field(default_factory=list)