diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 9f98c53b4cc..c77a3fb9ead 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -21,6 +21,7 @@ ) from miles.backends.training_utils.loss import compute_advantages_and_returns, get_log_probs_and_entropy, loss_function from miles.backends.training_utils.parallel import get_parallel_state, set_parallel_state +from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks from miles.ray.train_actor import TrainRayActor from miles.utils import async_utils, train_dump_utils, train_metric_utils from miles.utils.context_utils import with_defer @@ -32,6 +33,7 @@ from miles.utils.processing_utils import load_processor, load_tokenizer from miles.utils.profile_utils import TrainProfiler from miles.utils.ray_utils import Box +from miles.utils.sampling_mask import sampling_support_replay_enabled from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils.tracking import init_tracking @@ -369,6 +371,7 @@ def _compute_log_prob( active_model.eval() else: active_model = self.model + use_rollout_sampling_mask = model_tag == "actor" and sampling_support_replay_enabled(self.args) try: forward_data_store = [] @@ -392,6 +395,8 @@ def _compute_log_prob( "response_lengths", "max_seq_lens", ] + if use_rollout_sampling_mask: + forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"]) batch = get_batch( data_iterator, forward_only_keys, @@ -399,6 +404,9 @@ def _compute_log_prob( self.args.qkv_format, get_position_ids=True, ) + rollout_sampling_mask = None + if use_rollout_sampling_mask: + rollout_sampling_mask = get_rollout_sampling_masks(batch) model_args = self._get_model_inputs_args(batch) # keep logits in native bf16 (chunks upcast to fp32 downstream); avoids a full-vocab fp32 tensor (~5GB) @@ -413,6 +421,7 @@ def _compute_log_prob( response_lengths=batch["response_lengths"], with_entropy=(store_prefix == ""), max_seq_lens=batch.get("max_seq_lens", None), + rollout_sampling_mask=rollout_sampling_mask, ) batch_result = { @@ -501,6 +510,11 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: with routing_replay.stage(routing_replay.REPLAY_BACKWARD), timer("actor_train"): data_iterator.reset() num_steps_per_rollout = len(num_microbatches) + sampling_mask_keys = ( + ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets") + if sampling_support_replay_enabled(self.args) + else () + ) for step_id in range(num_steps_per_rollout): self.optimizer.zero_grad(set_to_none=True) @@ -523,6 +537,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: "returns", "ref_log_probs", "rollout_log_probs", + *sampling_mask_keys, ], self.args.data_pad_size_multiplier, self.args.qkv_format, diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index ee9cd157652..c65d001d8d9 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -33,6 +33,7 @@ from miles.utils.ray_utils import Box from miles.utils.reloadable_process_group import destroy_process_groups, monkey_patch_torch_dist, reload_process_groups from miles.utils.replay_base import all_replay_managers, routing_replay_manager +from miles.utils.sampling_mask import sampling_support_replay_enabled from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.timer import Timer, inverse_timer, timer from miles.utils.tracking_utils.structured_log import with_logs @@ -430,6 +431,7 @@ def compute_log_prob( rollout_id=rollout_id, store_prefix=store_prefix, fp32_output=False, + use_rollout_sampling_mask=store_prefix == "" and sampling_support_replay_enabled(self.args), ) @with_logs diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index aff3307f551..20182e968b6 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -29,11 +29,13 @@ from miles.backends.megatron_utils.ft.types import TrainStepOutcome from miles.backends.megatron_utils.local_weight_checksum import dump_local_weight_checksums from miles.backends.megatron_utils.optimizer_state_reset import reset_optimizer_states +from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks from miles.utils.audit_utils.witness.allocator import WitnessInfo from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase from miles.utils.memory_utils import clear_memory from miles.utils.multi_lora import is_multi_lora_enabled +from miles.utils.sampling_mask import sampling_support_replay_enabled from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.tracking_utils.structured_log import log_structured @@ -268,6 +270,7 @@ def forward_only( rollout_id: int, store_prefix: str = "", fp32_output: bool = True, + use_rollout_sampling_mask: bool = False, ) -> dict[str, list[torch.Tensor]]: """Run forward passes only and collect non-loss outputs (e.g., logprobs). @@ -283,6 +286,8 @@ def forward_only( rollout_id: Rollout identifier (selects the per-rollout dump subdirectory). store_prefix: Prefix to prepend to stored output keys. fp32_output: Whether Megatron should upcast the complete model output to FP32. + use_rollout_sampling_mask: Whether to score over each rollout token's + captured sampling support. Returns: Aggregated outputs keyed by ``store_prefix + key``. @@ -304,17 +309,20 @@ def forward_step( ) -> tuple[torch.Tensor, Callable[[torch.Tensor], dict[str, list[torch.Tensor]]]]: """Return the model output and its batch-bound loss callback.""" + forward_only_keys = [ + "tokens", + "loss_masks", + "multimodal_train_inputs", + "total_lengths", + "response_lengths", + "max_seq_lens", + "witness_ids", + ] + if use_rollout_sampling_mask: + forward_only_keys.extend(["rollout_sampling_mask_ids", "rollout_sampling_mask_offsets"]) batch = get_batch( data_iterator, - [ - "tokens", - "loss_masks", - "multimodal_train_inputs", - "total_lengths", - "response_lengths", - "max_seq_lens", - "witness_ids", - ], + forward_only_keys, args.data_pad_size_multiplier, args.qkv_format, allgather_cp=args.allgather_cp, @@ -324,6 +332,7 @@ def forward_step( packed_seq_params = get_packed_seq_params(batch, args) total_lengths = batch["total_lengths"] response_lengths = batch["response_lengths"] + rollout_sampling_mask = get_rollout_sampling_masks(batch) if use_rollout_sampling_mask else None if "adapter_token_counts" in batch: from megatron.bridge.peft.multi_lora_layers import set_tokens_per_adapter_slot @@ -342,8 +351,7 @@ def forward_step( fp32_output=fp32_output, ) - return output_tensor, partial( - f, + callback_kwargs = dict( args=args, unconcat_tokens=unconcat_tokens, total_lengths=total_lengths, @@ -351,6 +359,10 @@ def forward_step( with_entropy=args.use_rollout_entropy, max_seq_lens=batch.get("max_seq_lens", None), ) + if use_rollout_sampling_mask: + callback_kwargs["rollout_sampling_mask"] = rollout_sampling_mask + + return output_tensor, partial(f, **callback_kwargs) # Turn on evaluation mode which disables dropout. for model_module in model: @@ -407,6 +419,12 @@ def run_forward_backward_pass( ): """One pipeline forward/backward pass over the microbatches; no optimizer interaction.""" + sampling_mask_keys = ( + ("rollout_sampling_mask_ids", "rollout_sampling_mask_offsets") + if sampling_support_replay_enabled(args) and args.loss_type == "policy_loss" + else () + ) + @dumper_phase_util.wrap_forward_step def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_plan: bool = False) -> tuple[ torch.Tensor, @@ -447,6 +465,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "loss_weights", "target_tokens", "sample_indices", + *sampling_mask_keys, ], args.data_pad_size_multiplier, args.qkv_format, diff --git a/miles/backends/training_utils/loss_hub/losses.py b/miles/backends/training_utils/loss_hub/losses.py index d2d40b06936..60926e17bde 100644 --- a/miles/backends/training_utils/loss_hub/losses.py +++ b/miles/backends/training_utils/loss_hub/losses.py @@ -19,7 +19,9 @@ compute_policy_loss, ) from miles.backends.training_utils.parallel import get_parallel_state +from miles.backends.training_utils.sampling_mask import get_rollout_sampling_masks from miles.utils.function_registry import load_function +from miles.utils.sampling_mask import sampling_support_replay_enabled from miles.utils.types import RolloutBatch @@ -111,6 +113,7 @@ def policy_loss_function( total_lengths = batch["total_lengths"] max_seq_lens = batch.get("max_seq_lens", None) calculate_entropy = args.entropy_coef != 0 or args.observe_training_entropy + rollout_sampling_mask = get_rollout_sampling_masks(batch) if sampling_support_replay_enabled(args) else None log_probs_and_entropy = get_log_probs_and_entropy( logits, @@ -121,6 +124,7 @@ def policy_loss_function( with_entropy=calculate_entropy, entropy_requires_grad=args.entropy_coef != 0, max_seq_lens=max_seq_lens, + rollout_sampling_mask=rollout_sampling_mask, ) log_probs = log_probs_and_entropy["log_probs"] diff --git a/miles/backends/training_utils/sampling_mask.py b/miles/backends/training_utils/sampling_mask.py index a189f13a532..ca285983d76 100644 --- a/miles/backends/training_utils/sampling_mask.py +++ b/miles/backends/training_utils/sampling_mask.py @@ -1,10 +1,26 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence import torch from miles.utils.sampling_mask import RolloutSamplingMask +def get_rollout_sampling_masks(batch: Mapping[str, object]) -> list[RolloutSamplingMask]: + """Reconstruct the per-sample masks carried by the training-data wire.""" + ids_batch = batch.get("rollout_sampling_mask_ids") + offsets_batch = batch.get("rollout_sampling_mask_offsets") + if ids_batch is None or offsets_batch is None: + raise ValueError("sampling-support actor scoring requires both sampling-mask wire fields") + if not isinstance(ids_batch, Sequence) or not isinstance(offsets_batch, Sequence): + raise TypeError("rollout sampling-mask ids and offsets must be sequences with one entry per sample") + if len(ids_batch) != len(offsets_batch): + raise ValueError(f"sampling-mask ids batch size {len(ids_batch)} != offsets batch size {len(offsets_batch)}") + return [ + RolloutSamplingMask(ids=torch.as_tensor(ids), offsets=torch.as_tensor(offsets)) + for ids, offsets in zip(ids_batch, offsets_batch, strict=True) + ] + + def build_local_sampling_mask( logits: torch.Tensor, sampling_mask: RolloutSamplingMask, diff --git a/miles/utils/sampling_mask.py b/miles/utils/sampling_mask.py index 568d0ebb396..164abf60684 100644 --- a/miles/utils/sampling_mask.py +++ b/miles/utils/sampling_mask.py @@ -5,6 +5,11 @@ import torch +def sampling_support_replay_enabled(args) -> bool: + """Whether rollout sampling uses a bounded support that must be replayed.""" + return float(getattr(args, "rollout_top_p", 1.0)) < 1.0 or int(getattr(args, "rollout_top_k", -1)) > 0 + + @dataclass(frozen=True, eq=False) class RolloutSamplingMask: """One sample's sampling mask: for each response position, the token ids diff --git a/tests/fast/backends/megatron_utils/test_model.py b/tests/fast/backends/megatron_utils/test_model.py index cea3a0fa231..c3e69f68c02 100644 --- a/tests/fast/backends/megatron_utils/test_model.py +++ b/tests/fast/backends/megatron_utils/test_model.py @@ -2,6 +2,7 @@ from argparse import Namespace from dataclasses import dataclass, field from typing import Any +from unittest.mock import Mock import pytest @@ -120,3 +121,57 @@ def test_train_one_step_emits_the_train_tag_in_its_structured_event( ) assert "train op=train_step rollout=7 step=3 attempt=2 outcome=NORMAL valid_step=true" in caplog.messages + + +def test_forward_only_omits_sampling_mask_for_callbacks_that_do_not_replay_sampling_support(monkeypatch): + from miles.backends.megatron_utils import model as model_module + + def forward_backward(**kwargs): + output, callback = kwargs["forward_step_func"](kwargs["data_iterator"][0], kwargs["model"][0]) + return [callback(output)] + + callback_kwargs = {} + + def critic_callback(_output, **kwargs): + callback_kwargs.update(kwargs) + return {} + + args = Namespace( + data_pad_size_multiplier=1, + qkv_format="thd", + allgather_cp=False, + enable_witness=False, + use_rollout_entropy=False, + custom_megatron_before_log_prob_hook_path=None, + seq_length=8, + micro_batch_size=1, + ) + batch = { + "tokens": None, + "input_loss_masks": None, + "multimodal_train_inputs": None, + "unconcat_tokens": [], + "total_lengths": [], + "response_lengths": [], + "max_seq_lens": None, + } + dumper = Mock() + dumper.wrap_forward_step.side_effect = lambda fn: fn + + monkeypatch.setattr(model_module, "DumperMegatronUtil", lambda *_args, **_kwargs: dumper) + monkeypatch.setattr(model_module, "get_model_config", lambda _model: Mock(timers=None)) + monkeypatch.setattr(model_module, "get_batch", lambda *_args, **_kwargs: batch) + monkeypatch.setattr(model_module, "get_packed_seq_params", lambda *_args, **_kwargs: None) + monkeypatch.setattr(model_module, "get_forward_backward_func", lambda: forward_backward) + monkeypatch.setattr(model_module, "mpu", FakeMpu()) + + model_module.forward_only( + critic_callback, + args, + [Mock()], + [Mock()], + [1], + rollout_id=0, + ) + + assert "rollout_sampling_mask" not in callback_kwargs diff --git a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py index caa92f15b1a..aa80bb54f06 100644 --- a/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py +++ b/tests/fast/backends/megatron_utils/test_shared_ppo_lifecycle.py @@ -136,18 +136,25 @@ def test_train_keeps_model_resident(actor_module, monkeypatch): worker.sleep.assert_not_called() -def test_compute_log_prob_keeps_logits_in_model_precision(actor_module, monkeypatch): +@pytest.mark.parametrize( + ("store_prefix", "expected"), + [("", True), ("ref_", False), ("teacher_", False)], +) +def test_compute_log_prob_replays_sampling_support_only_for_actor_scores( + actor_module, monkeypatch, store_prefix, expected +): worker = object.__new__(actor_module.MegatronTrainRayActor) - worker.args = Namespace() + worker.args = Namespace(rollout_top_p=0.95, rollout_top_k=32) worker.model = [object()] forward_only = Mock(return_value={"log_probs": []}) monkeypatch.setattr(actor_module, "forward_only", forward_only) monkeypatch.setattr(actor_module, "timer", lambda _name: nullcontext()) - result = worker.compute_log_prob([], [], rollout_id=3) + result = worker.compute_log_prob([], [], rollout_id=3, store_prefix=store_prefix) assert result == {"log_probs": []} assert forward_only.call_args.kwargs["fp32_output"] is False + assert forward_only.call_args.kwargs["use_rollout_sampling_mask"] is expected def test_save_model_does_not_manage_lifecycle(actor_module, monkeypatch): diff --git a/tests/fast/backends/test_fsdp_train.py b/tests/fast/backends/test_fsdp_train.py index 9191bd9a37a..646f8e8354e 100644 --- a/tests/fast/backends/test_fsdp_train.py +++ b/tests/fast/backends/test_fsdp_train.py @@ -1,8 +1,12 @@ from argparse import Namespace from collections.abc import Iterator from contextlib import contextmanager, nullcontext +from types import SimpleNamespace from unittest.mock import Mock +import pytest +import torch + from miles.backends.fsdp_utils import actor as actor_module from miles.backends.megatron_utils.ft.types import TrainStepOutcome, TrainStepOutput @@ -29,3 +33,57 @@ def test_fsdp_train_debug_rollout_only_returns_a_normal_output(monkeypatch): assert result == TrainStepOutput(outcome=TrainStepOutcome.NORMAL) actor._train_core.assert_not_called() + + +@pytest.mark.parametrize( + ("model_tag", "store_prefix", "top_p", "top_k", "expected"), + [ + ("actor", "", 0.95, 32, True), + ("ref", "ref_", 0.95, 32, False), + ("actor", "", 1.0, -1, False), + ], +) +def test_fsdp_only_actor_scores_replay_rollout_sampling_support( + monkeypatch, model_tag, store_prefix, top_p, top_k, expected +): + actor = object.__new__(actor_module.FSDPTrainRayActor) + actor.args = Namespace( + rollout_top_p=top_p, + rollout_top_k=top_k, + data_pad_size_multiplier=1, + qkv_format="thd", + ) + actor.model = Mock(return_value=SimpleNamespace(logits=torch.zeros(1, 1, 4))) + actor.ref_model = None + actor.precision_policy = object() + actor.prof = SimpleNamespace(iterate_train_log_probs=lambda values: values) + actor._get_model_inputs_args = lambda _batch: {} + data_iterator = Mock() + sampling_masks = [object()] + captured = {} + + monkeypatch.setattr(actor_module, "timer", _noop_timer) + monkeypatch.setattr(actor_module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(actor_module, "precision_forward_context", lambda _policy: nullcontext()) + monkeypatch.setattr( + actor_module, + "get_batch", + lambda *_args, **_kwargs: { + "unconcat_tokens": [torch.tensor([0])], + "total_lengths": [1], + "response_lengths": [1], + "max_seq_lens": None, + }, + ) + monkeypatch.setattr(actor_module, "get_rollout_sampling_masks", lambda _batch: sampling_masks) + + def get_log_probs_and_entropy(**kwargs): + captured.update(kwargs) + return {"log_probs": [torch.zeros(1)]} + + monkeypatch.setattr(actor_module, "get_log_probs_and_entropy", get_log_probs_and_entropy) + monkeypatch.setattr(actor_module, "aggregate_forward_results", lambda *_args, **_kwargs: {}) + + actor._compute_log_prob(model_tag, data_iterator, [1], store_prefix=store_prefix) + + assert (captured["rollout_sampling_mask"] is sampling_masks) is expected diff --git a/tests/fast/backends/training_utils/test_sampling_mask.py b/tests/fast/backends/training_utils/test_sampling_mask.py index 0a70dce860f..47b90a47b75 100644 --- a/tests/fast/backends/training_utils/test_sampling_mask.py +++ b/tests/fast/backends/training_utils/test_sampling_mask.py @@ -6,10 +6,28 @@ from miles.backends.training_utils import cp_utils from miles.backends.training_utils.loss_hub import logit_processors from miles.backends.training_utils.loss_hub.math_utils import _calculate_log_probs_and_entropy_true_on_policy -from miles.backends.training_utils.sampling_mask import build_local_sampling_mask +from miles.backends.training_utils.sampling_mask import build_local_sampling_mask, get_rollout_sampling_masks from miles.utils.sampling_mask import RolloutSamplingMask +def test_get_rollout_sampling_masks_reconstructs_transport_batch(): + masks = get_rollout_sampling_masks( + { + "rollout_sampling_mask_ids": [torch.tensor([1, 4, 2], dtype=torch.int32)], + "rollout_sampling_mask_offsets": [torch.tensor([0, 2, 3], dtype=torch.int64)], + } + ) + + ids, offsets = masks[0]._as_tensors() + assert ids.tolist() == [1, 4, 2] + assert offsets.tolist() == [0, 2, 3] + + +def test_get_rollout_sampling_masks_requires_complete_transport_batch(): + with pytest.raises(ValueError, match="requires both sampling-mask wire fields"): + get_rollout_sampling_masks({"rollout_sampling_mask_ids": [torch.tensor([1])]}) + + def test_build_local_sampling_mask_selects_original_response_rows_and_tp_shard(): logits = torch.zeros(2, 4) mask = build_local_sampling_mask( @@ -101,6 +119,31 @@ def test_true_on_policy_masks_logprob_but_keeps_full_vocab_entropy(): torch.testing.assert_close(entropy, expected_entropy) +def test_sampling_support_normalization_has_the_expected_actor_gradient(): + logits = torch.tensor( + [[2.0, 1.0, 0.0, -1.0], [-1.0, 2.0, 1.0, 0.0]], + dtype=torch.float64, + requires_grad=True, + ) + tokens = torch.tensor([0, 1]) + sampling_mask = torch.tensor([[True, False, True, False], [False, True, False, False]]) + + log_probs, _ = _calculate_log_probs_and_entropy_true_on_policy( + logits, + tokens, + None, + sampling_mask=sampling_mask, + ) + log_probs.sum().backward() + + first_support_probs = torch.softmax(logits.detach()[0, [0, 2]], dim=-1) + expected_grad = torch.zeros_like(logits) + expected_grad[0, 0] = 1 - first_support_probs[0] + expected_grad[0, 2] = -first_support_probs[1] + + torch.testing.assert_close(logits.grad, expected_grad) + + def test_get_log_probs_and_entropy_applies_per_response_sampling_support(monkeypatch): parallel_state = SimpleNamespace( tp=SimpleNamespace(rank=0, group=None),