From 4d17e8a081c42f29ab74db8dcbba4ccab066581f Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 3 Jun 2026 01:05:05 +0000 Subject: [PATCH 01/28] add mtp support --- .../skyrl_train/inference_servers/utils.py | 12 ++- .../backends/skyrl_train/utils/torch_utils.py | 38 +++++++ .../megatron/megatron_model_wrapper.py | 99 ++++++++++++++++++- .../workers/megatron/megatron_worker.py | 28 ++++++ .../workers/megatron/model_bridges.py | 10 +- skyrl/train/config/config.py | 22 +++++ .../test_build_vllm_cli_args.py | 20 ++++ .../skyrl_train/utils/test_mtp_torch_utils.py | 52 ++++++++++ tests/train/test_mtp_config.py | 37 +++++++ 9 files changed, 313 insertions(+), 5 deletions(-) create mode 100644 tests/backends/skyrl_train/utils/test_mtp_torch_utils.py create mode 100644 tests/train/test_mtp_config.py diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index 9705b0e4b6..34231df827 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -143,7 +143,17 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: else: args.enable_lora = False - # Add any extra engine_init_kwargs + # Speculative decoding (e.g. Multi-Token Prediction). Pass the dict straight through to vLLM's + # AsyncEngineArgs.speculative_config. For MTP (`{"method": "mtp", ...}`) vLLM loads the MTP heads + # from the same policy checkpoint, so SkyRL's weight sync keeps the draft aligned with the policy + # (the trained MTP head weights are exported by the Megatron bridge and loaded into the draft). + if ie_cfg.speculative_config is not None: + spec_cfg = get_config_as_dict(ie_cfg.speculative_config) + args.speculative_config = spec_cfg + logger.info(f"vLLM speculative decoding enabled: speculative_config={spec_cfg}") + + # Add any extra engine_init_kwargs (applied last so they can override anything above, including + # speculative_config, if a user needs full control). engine_kwargs = get_config_as_dict(ie_cfg.engine_init_kwargs) for key, value in engine_kwargs.items(): setattr(args, key, value) diff --git a/skyrl/backends/skyrl_train/utils/torch_utils.py b/skyrl/backends/skyrl_train/utils/torch_utils.py index 3da9dc4077..2ac21291c2 100644 --- a/skyrl/backends/skyrl_train/utils/torch_utils.py +++ b/skyrl/backends/skyrl_train/utils/torch_utils.py @@ -184,6 +184,44 @@ def masked_mean(tensor: torch.Tensor, mask: torch.Tensor | None, dim: int | None return (tensor * mask).sum(axis=dim) / mask.sum(axis=dim).clamp(min=1.0) +def build_mtp_next_token_labels(sequences: torch.Tensor) -> torch.Tensor: + """Build pre-shifted next-token labels for Multi-Token Prediction (MTP). + + Megatron's MTP path (``process_mtp_loss``) expects ``labels`` in the same convention as + the main language-model loss: ``labels[t]`` is the token the model should predict at + position ``t`` (i.e. ``sequences[t + 1]``). Megatron then rolls these labels once more per + MTP layer so MTP layer ``k`` is supervised against token ``t + k + 2``. + + The last position has no next token, so it is zeroed here; the corresponding loss-mask + entry is dropped by Megatron's own boundary handling (``roll_tensor`` zeros the wrapped + position), so it never contributes to the loss. + + Args: + sequences: ``[batch, seq_len]`` token ids (per-row a single, contiguous sequence). + + Returns: + ``[batch, seq_len]`` next-token labels (same dtype as ``sequences``). + """ + labels = torch.roll(sequences, shifts=-1, dims=1) + labels[:, -1] = 0 + return labels + + +def build_mtp_loss_mask(attention_mask: torch.Tensor) -> torch.Tensor: + """Build the MTP loss mask from an attention mask. + + We train the MTP heads on every real token of the sequence (next-token prediction over the + full prompt+response), mirroring standard MTP pretraining. ``attention_mask`` already marks + real (non-pad) tokens, so it doubles as the loss mask. Per-sequence boundaries (the final + real token, whose MTP target is invalid) are handled by Megatron's roll in + ``process_mtp_loss`` and need no special treatment here. + + Returns a float tensor so it can flow through the same de-padding/packing transforms applied + to the token ids. + """ + return attention_mask.to(torch.float32) + + def safe_exp_delta(delta: torch.Tensor, clip: float = 20.0, out_dtype=None) -> torch.Tensor: """ Clamp the delta before exponentiating to avoid potential overflow. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index ac7258fac7..e70836743b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -29,10 +29,42 @@ setup_per_microbatch_replay_backward, setup_per_microbatch_replay_forward, ) -from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean +from skyrl.backends.skyrl_train.utils.torch_utils import ( + build_mtp_loss_mask, + build_mtp_next_token_labels, + masked_mean, +) from skyrl.train.config import TrainerConfig +def _mtp_logits_passthrough( + *, + hidden_states, + output_layer, + output_weight, + runtime_gather_output=None, + scale_logits=None, + **kwargs, +): + """``output_processor`` hook for GPTModel that returns logits instead of the LM loss. + + When Multi-Token Prediction is active we pass ``labels``/``loss_mask`` into ``GPTModel.forward`` + so that ``process_mtp_loss`` runs and injects the MTP loss into the autograd graph (via + ``MTPLossAutoScaler`` on ``hidden_states``). But passing labels would normally make the model + return the *main* LM loss, whereas SkyRL needs the *logits* to compute its own PPO/GRPO/CE loss. + + GPTModel invokes this hook after ``process_mtp_loss`` and returns whatever it produces, short- + circuiting the labels->loss path. We reproduce the stock ``labels is None`` logits path exactly + (see ``GPTModel._postprocess``) so all downstream SkyRL logic is unchanged: vocab-parallel + logits in ``[batch, seq, vocab]`` layout. + """ + logits, _ = output_layer(hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output) + if scale_logits is not None: + logits = scale_logits(logits) + # [s b h] => [b s h], matching GPTModel's non-loss return. + return logits.transpose(0, 1).contiguous() + + class MegatronModelWrapper: def __init__( self, @@ -226,6 +258,32 @@ def forward_backward_mini_batch( """ forward_backward_func = get_forward_backward_func() + # ------------------------------------------------------------------ + # Multi-Token Prediction (MTP) + # ------------------------------------------------------------------ + # If the model was built with MTP heads (policy.megatron_config.mtp_num_layers), train them + # alongside the policy by feeding next-token labels into GPTModel.forward so Megatron's + # process_mtp_loss runs and injects the MTP gradient via MTPLossAutoScaler. We only do this + # when actually training (not forward_only / logprob-only passes). + model_config = get_model_config(self.actor_module[0]) + mtp_enabled = (not forward_only) and bool(getattr(model_config, "mtp_num_layers", None)) + if mtp_enabled: + from megatron.core.transformer.multi_token_prediction import ( + MTPLossAutoScaler, + MTPLossLoggingHelper, + ) + + # Match the pipeline schedule's per-microbatch loss division (Megatron divides the + # main loss by num_microbatches before backward). The MTP autoscaler injects a constant + # gradient independent of the main loss value, so we scale it the same way to keep the + # MTP aux-loss gradient commensurate with the policy-loss gradient. mtp_loss_scaling_factor + # (set on the provider) is the higher-level knob for the MTP/policy loss ratio. + MTPLossAutoScaler.set_loss_scale( + torch.tensor(1.0 / max(1, len(micro_batches)), device=torch.cuda.current_device()) + ) + if "values" in MTPLossLoggingHelper.tracker: + MTPLossLoggingHelper.clean_loss_in_tracker() + # Resolve loss function resolved_loss_name = loss_fn if loss_fn is not None else self.cfg.algorithm.policy_loss_type if loss_fn is not None: @@ -440,11 +498,34 @@ def forward_step(batch_iter, model): ) packed_seq_params = None + # When MTP is active, build next-token labels + loss mask in the SAME de-padded/packed + # space as new_sequences and pass them (plus the logits passthrough) so process_mtp_loss + # runs. We materialize labels on all ranks (pre_process=True) because they are consumed + # on the post_process / last pipeline stage, regardless of which stage owns the input. + mtp_kwargs = {} + if mtp_enabled: + labels_full = build_mtp_next_token_labels(sequences) + loss_mask_full = build_mtp_loss_mask(attention_mask) + if self.remove_microbatch_padding: + mtp_labels = preprocess_packed_seqs(labels_full, attention_mask, pre_process=True)[0] + mtp_loss_mask = preprocess_packed_seqs(loss_mask_full, attention_mask, pre_process=True)[0] + else: + mtp_labels = remove_left_padding(labels_full, attention_mask, position_ids, pre_process=True)[0] + mtp_loss_mask = remove_left_padding(loss_mask_full, attention_mask, position_ids, pre_process=True)[ + 0 + ] + mtp_kwargs = dict( + labels=mtp_labels, + loss_mask=mtp_loss_mask, + output_processor=_mtp_logits_passthrough, + ) + outputs = model( new_sequences, new_position_ids, new_attention_mask, packed_seq_params=packed_seq_params, + **mtp_kwargs, ) if self.remove_microbatch_padding: @@ -483,6 +564,22 @@ def forward_step(batch_iter, model): forward_only=forward_only, ) + # Surface the MTP loss for logging. process_mtp_loss accumulates a per-layer loss into + # MTPLossLoggingHelper.tracker (summed across microbatches) on the last pipeline stage. + # Attach the mean MTP loss to the first microbatch's metrics so it rides along the existing + # PP broadcast + DP reduction below. + if mtp_enabled and mpu.is_pipeline_last_stage(ignore_virtual=True): + from megatron.core.transformer.multi_token_prediction import ( + MTPLossLoggingHelper, + ) + + mtp_values = MTPLossLoggingHelper.tracker.get("values") + if mtp_values is not None and metrics_list and metrics_list[0] is not None: + mtp_mean = (mtp_values.sum() / (mtp_values.numel() * max(1, len(micro_batches)))).item() + metrics_list[0]["mtp_loss"] = mtp_mean + if mtp_values is not None: + MTPLossLoggingHelper.clean_loss_in_tracker() + # broadcast metrics to all pp ranks if not mpu.is_pipeline_last_stage(ignore_virtual=True): metrics_list = [None] * len(micro_batches) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index b2257fb964..384dbac6e7 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -326,9 +326,16 @@ def init_configs( bf16=True, flash_attn=False, lora_config=None, + enable_mtp=True, ): """ Initialize the Megatron-Bridge bridge and provider objects + hf_config and tokenizer + + Args: + enable_mtp: When True (policy worker), honor Multi-Token Prediction (MTP) heads — + either from ``megatron_config.mtp_num_layers`` or, if that is None, from the + model's own HF config (``num_nextn_predict_layers``). When False (ref worker / + inference-only flows), force MTP off so the extra layers are neither built nor run. """ tokenizer = get_tokenizer(model_path, trust_remote_code=True) hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -397,6 +404,25 @@ def init_configs( # Apply any additional transformer config kwargs (can override the above). for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) + + # Multi-Token Prediction (MTP) head configuration. + # DeepSeek-V3 / GLM-4.7-Flash bridges set provider.mtp_num_layers from the HF config's + # num_nextn_predict_layers. We layer SkyRL's explicit config on top: + # - enable_mtp=False (ref worker): force MTP off (the heads don't affect ref logprobs + # and would only waste compute/memory). + # - mtp_num_layers set: override the model default (0 force-disables). + # - mtp_num_layers=None: keep whatever the bridge inferred from the model. + if not enable_mtp: + provider.mtp_num_layers = None + elif megatron_config.mtp_num_layers is not None: + provider.mtp_num_layers = megatron_config.mtp_num_layers or None + if getattr(provider, "mtp_num_layers", None): + provider.mtp_loss_scaling_factor = megatron_config.mtp_loss_scaling_factor + logger.info( + f"MTP enabled: mtp_num_layers={provider.mtp_num_layers}, " + f"mtp_loss_scaling_factor={provider.mtp_loss_scaling_factor}" + ) + provider.finalize() self.provider = provider @@ -1250,6 +1276,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): self.cfg.ref.megatron_config.transformer_config_kwargs, bf16=self.cfg.bf16, flash_attn=self.cfg.flash_attn, + # Ref worker only needs main-model logprobs; MTP heads are irrelevant here. + enable_mtp=False, ) self.actor_module = self.make_megatron_module( diff --git a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py index ee8e2160db..327334bc5e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py +++ b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py @@ -63,9 +63,13 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): hf_config.rope_theta = orig_rope_theta provider.moe_router_score_function = "sigmoid" - # TODO (erictang000): follow up when Megatron-Bridge supports MTP - # layers for DeepSeek-V3 style models - provider.mtp_num_layers = None + # NOTE: MTP is now honored. DeepSeekV3Bridge.provider_bridge already sets + # provider.mtp_num_layers from hf_config.num_nextn_predict_layers, and + # megatron-bridge's get_common_mapping_list emits the MTP weight mappings + # (enorm/hnorm/eh_proj/shared_head.* + nextn transformer layers) for HF + # round-tripping. SkyRL's MegatronWorker.init_configs can still override + # provider.mtp_num_layers via policy.megatron_config.mtp_num_layers (e.g. + # set it to 0 to force-disable, or leave None to use the model default). return provider except ImportError: diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 60a6f4d0a8..ea730df8a9 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -194,6 +194,21 @@ class MegatronConfig(BaseConfig): freeze_moe_router: bool = False """If True, freeze MoE router parameters so they are not updated during training. No-op on non-MoE models.""" + mtp_num_layers: Optional[int] = None + """Number of Multi-Token Prediction (MTP) layers/heads to build on the model. + + - ``None`` (default): honor the model's own HF config (``num_nextn_predict_layers``). For models + that ship MTP heads (e.g. GLM-4.7-Flash / DeepSeek-V3) this enables and trains them; for models + without MTP heads it is a no-op. + - An ``int``: explicitly override (e.g. ``0`` to force-disable MTP even for MTP-capable models). + + When MTP layers are present, SkyRL adds the Megatron MTP loss to the training step so the heads + track the updated policy, and the heads are synced to vLLM for speculative decoding (set + ``generator.inference_engine.speculative_config`` accordingly).""" + mtp_loss_scaling_factor: float = 0.1 + """Scaling coefficient for the MTP auxiliary loss (Megatron ``mtp_loss_scaling_factor``). + Only used when MTP layers are active. Higher values train the MTP heads more aggressively at the + cost of pulling on the shared trunk; the default matches Megatron/DeepSeek-V3.""" def __post_init__(self): # Backfill defaults for any keys the user didn't override so an override dict @@ -525,6 +540,13 @@ class InferenceEngineConfig(BaseConfig): multimodal models (e.g. Qwen3.5) skip vision encoder initialization.""" engine_init_kwargs: Dict[str, Any] = field(default_factory=dict) """Pass-through kwargs for the vLLM engine. Names must match the engine's args.""" + speculative_config: Optional[Dict[str, Any]] = None + """Speculative-decoding config passed through to vLLM (``AsyncEngineArgs.speculative_config``). + Use this to enable Multi-Token Prediction (MTP) draft decoding for faster rollout, e.g. + ``{"method": "mtp", "num_speculative_tokens": 1}``. With ``method="mtp"`` vLLM loads the MTP + heads from the *same* policy checkpoint, so SkyRL's weight sync keeps the draft in sync with + the trained policy (requires ``policy.megatron_config.mtp_num_layers`` > 0 on the training side + for the heads to actually be trained). ``None`` disables speculative decoding.""" override_existing_update_group: str = "auto" """``"auto"``, ``"enable"``, or ``"disable"``.""" external_proxy_url: Optional[str] = None diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py index b386e6cdd7..650a3464c2 100644 --- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py +++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py @@ -26,3 +26,23 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): assert args.model == cfg.trainer.policy.model.path assert args.tensor_parallel_size == cfg.generator.inference_engine.tensor_parallel_size assert vllm.platforms.current_platform.device_type == "cuda" + + +@pytest.mark.vllm +def test_build_vllm_cli_args_speculative_config_mtp(monkeypatch): + """speculative_config is passed through to vLLM for MTP draft decoding.""" + import vllm.platforms + from vllm.platforms.interface import UnspecifiedPlatform + + monkeypatch.setattr(vllm.platforms, "_current_platform", UnspecifiedPlatform()) + + cfg = SkyRLTrainConfig() + # Default: no speculative decoding. + args = build_vllm_cli_args(cfg) + assert getattr(args, "speculative_config", None) is None + + # Enable MTP speculative decoding. + spec = {"method": "mtp", "num_speculative_tokens": 1} + cfg.generator.inference_engine.speculative_config = spec + args = build_vllm_cli_args(cfg) + assert args.speculative_config == spec diff --git a/tests/backends/skyrl_train/utils/test_mtp_torch_utils.py b/tests/backends/skyrl_train/utils/test_mtp_torch_utils.py new file mode 100644 index 0000000000..dee9a9068d --- /dev/null +++ b/tests/backends/skyrl_train/utils/test_mtp_torch_utils.py @@ -0,0 +1,52 @@ +"""CPU unit tests for the MTP (Multi-Token Prediction) label helpers. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/utils/test_mtp_torch_utils.py +""" + +import torch + +from skyrl.backends.skyrl_train.utils.torch_utils import ( + build_mtp_loss_mask, + build_mtp_next_token_labels, +) + + +def test_build_mtp_next_token_labels_basic(): + # labels[t] should be sequences[t + 1]; last position zeroed (no next token). + sequences = torch.tensor([[10, 11, 12, 13], [20, 21, 22, 23]]) + labels = build_mtp_next_token_labels(sequences) + expected = torch.tensor([[11, 12, 13, 0], [21, 22, 23, 0]]) + assert torch.equal(labels, expected) + # dtype preserved (token ids) + assert labels.dtype == sequences.dtype + + +def test_build_mtp_next_token_labels_does_not_mutate_input(): + sequences = torch.tensor([[1, 2, 3]]) + seq_copy = sequences.clone() + _ = build_mtp_next_token_labels(sequences) + assert torch.equal(sequences, seq_copy) + + +def test_build_mtp_next_token_labels_no_wraparound_into_first_token(): + # The roll's wrapped value (original first token) must be overwritten with 0, + # so the model never sees the start token as a "next token" target at the end. + sequences = torch.tensor([[7, 8, 9]]) + labels = build_mtp_next_token_labels(sequences) + assert labels[0, -1].item() == 0 + assert labels[0, -1].item() != sequences[0, 0].item() or sequences[0, 0].item() == 0 + + +def test_build_mtp_loss_mask_from_attention_mask(): + attention_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 0, 0]]) + loss_mask = build_mtp_loss_mask(attention_mask) + assert loss_mask.dtype == torch.float32 + assert torch.equal(loss_mask, attention_mask.to(torch.float32)) + + +def test_build_mtp_loss_mask_accepts_bool_mask(): + # forward_step passes attention_mask.to(bool); the helper must still produce a float mask. + attention_mask = torch.tensor([[True, True, False]]) + loss_mask = build_mtp_loss_mask(attention_mask) + assert loss_mask.dtype == torch.float32 + assert torch.equal(loss_mask, torch.tensor([[1.0, 1.0, 0.0]])) diff --git a/tests/train/test_mtp_config.py b/tests/train/test_mtp_config.py new file mode 100644 index 0000000000..47469c724d --- /dev/null +++ b/tests/train/test_mtp_config.py @@ -0,0 +1,37 @@ +"""CPU unit tests for the MTP config knobs. + +uv run --isolated --extra dev pytest tests/train/test_mtp_config.py +""" + +from skyrl.train.config import InferenceEngineConfig, MegatronConfig +from skyrl.train.config.config import build_nested_dataclass + + +def test_megatron_config_mtp_defaults(): + cfg = MegatronConfig() + # None => honor the model's own num_nextn_predict_layers (no SkyRL override). + assert cfg.mtp_num_layers is None + assert cfg.mtp_loss_scaling_factor == 0.1 + + +def test_megatron_config_mtp_overrides_parse(): + cfg = build_nested_dataclass(MegatronConfig, {"mtp_num_layers": 2, "mtp_loss_scaling_factor": 0.3}) + assert cfg.mtp_num_layers == 2 + assert cfg.mtp_loss_scaling_factor == 0.3 + + +def test_megatron_config_mtp_force_disable(): + # An explicit 0 is how a user force-disables MTP even on an MTP-capable model. + cfg = build_nested_dataclass(MegatronConfig, {"mtp_num_layers": 0}) + assert cfg.mtp_num_layers == 0 + + +def test_inference_engine_speculative_config_default_none(): + cfg = InferenceEngineConfig() + assert cfg.speculative_config is None + + +def test_inference_engine_speculative_config_parses_mtp_dict(): + spec = {"method": "mtp", "num_speculative_tokens": 1} + cfg = build_nested_dataclass(InferenceEngineConfig, {"speculative_config": spec}) + assert cfg.speculative_config == spec From 1a9d7911f3ad336f52bf55dec7204df387e02e4e Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Fri, 5 Jun 2026 20:41:22 +0000 Subject: [PATCH 02/28] decoupled CE loss for MTP weight update --- skyrl/backends/skyrl_train/mtp/__init__.py | 32 +++ skyrl/backends/skyrl_train/mtp/adapter.py | 56 +++++ .../skyrl_train/mtp/draft_loss_wrapper.py | 118 +++++++++ .../skyrl_train/mtp/hidden_capture.py | 116 +++++++++ skyrl/backends/skyrl_train/mtp/soft_ce.py | 207 ++++++++++++++++ .../megatron/megatron_model_wrapper.py | 225 +++++++++--------- .../workers/megatron/megatron_worker.py | 12 +- skyrl/train/config/config.py | 31 ++- .../skyrl_train/mtp/test_hidden_capture.py | 86 +++++++ .../backends/skyrl_train/mtp/test_soft_ce.py | 116 +++++++++ tests/train/test_mtp_config.py | 21 +- 11 files changed, 904 insertions(+), 116 deletions(-) create mode 100644 skyrl/backends/skyrl_train/mtp/__init__.py create mode 100644 skyrl/backends/skyrl_train/mtp/adapter.py create mode 100644 skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py create mode 100644 skyrl/backends/skyrl_train/mtp/hidden_capture.py create mode 100644 skyrl/backends/skyrl_train/mtp/soft_ce.py create mode 100644 tests/backends/skyrl_train/mtp/test_hidden_capture.py create mode 100644 tests/backends/skyrl_train/mtp/test_soft_ce.py diff --git a/skyrl/backends/skyrl_train/mtp/__init__.py b/skyrl/backends/skyrl_train/mtp/__init__.py new file mode 100644 index 0000000000..0949ee010f --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/__init__.py @@ -0,0 +1,32 @@ +# Decoupled Multi-Token Prediction (MTP) draft-head training. +# +# This package mirrors the design used by NeMo-RL's draft-model training +# (https://github.com/NVIDIA-NeMo/RL, ``nemo_rl/models/megatron/draft`` and +# ``nemo_rl/algorithms/loss``): the trunk's hidden states are *detached* before +# the draft/MTP head runs (decoupling the draft gradient from the policy +# backbone), the head is supervised by an explicit, configurable loss +# (soft cross-entropy distillation against the policy's own next-token +# distribution, or hard next-token cross-entropy), and the two losses are +# combined legibly as ``policy_loss + w * draft_loss`` instead of being +# entangled inside Megatron's ``MTPLossAutoScaler``. +# +# The pieces here are backend-agnostic. Only the ``adapter`` wiring is +# backend-specific (Megatron today; FSDP is a planned follow-up). + +from skyrl.backends.skyrl_train.mtp.draft_loss_wrapper import ( + DraftLossWrapper, + combine_policy_and_draft_loss, +) +from skyrl.backends.skyrl_train.mtp.soft_ce import ( + build_teacher_logits, + draft_hard_ce, + draft_soft_ce, +) + +__all__ = [ + "DraftLossWrapper", + "combine_policy_and_draft_loss", + "build_teacher_logits", + "draft_hard_ce", + "draft_soft_ce", +] diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py new file mode 100644 index 0000000000..78c2eb39ff --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -0,0 +1,56 @@ +# Backend seam for decoupled MTP/draft-head training. +# +# The loss (``soft_ce``), the combination (``draft_loss_wrapper``) and the +# capture mechanism (``hidden_capture``) are backend-agnostic. Only two things +# differ per backend: (a) where the trunk hidden states are captured and (b) how +# the draft head turns those hidden states into vocab logits. This module pins +# down that contract. Megatron is implemented today; FSDP is a planned follow-up +# that will implement the same protocol and reuse everything else unchanged. + +from __future__ import annotations + +from typing import List, Protocol, runtime_checkable + + +@runtime_checkable +class DraftAdapter(Protocol): + """Per-backend hooks for decoupled MTP training.""" + + def capture_context(self, model, detach_trunk: bool): ... + + def project_to_logits(self, hidden_states_per_layer, model) -> List: ... + + +def project_mtp_hidden_to_logits(hidden_states_per_layer, gpt_model, detach_output_weight: bool = False) -> List: + """Run the shared output layer on each captured MTP hidden-state chunk. + + Mirrors ``GPTModel._postprocess``: logits come out of the output layer in + ``[seq, batch, vocab/tp]`` and are transposed to ``[batch, seq, vocab/tp]`` + to match the main-logits layout. The returned tensors are still in Megatron's + *internal* (packed / left-removed) sequence layout; the caller applies the + same de-padding transform used for the main logits so they align with the + teacher. + + Args: + hidden_states_per_layer: list of ``[seq, batch, hidden]`` tensors, one per + MTP depth (from ``MTPHiddenCapture.compute_student_hidden_states``). + gpt_model: the unwrapped ``GPTModel`` (exposes ``output_layer`` and, when + embeddings are tied, ``shared_embedding_or_output_weight``). + + Returns: + list of ``[batch, seq, vocab/tp]`` student-logits tensors (internal layout). + """ + output_weight = None + if getattr(gpt_model, "share_embeddings_and_output_weights", False): + output_weight = gpt_model.shared_embedding_or_output_weight() + if detach_output_weight: + # Optionally isolate the shared embedding/output-layer from the draft + # gradient too (default keeps NeMo's behaviour: only the trunk hidden + # states are detached). + output_weight = output_weight.detach() + + logits_per_layer = [] + for hidden in hidden_states_per_layer: + logits, _ = gpt_model.output_layer(hidden, weight=output_weight) + logits_per_layer.append(logits.transpose(0, 1).contiguous()) + return logits_per_layer diff --git a/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py b/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py new file mode 100644 index 0000000000..d44c7894fd --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py @@ -0,0 +1,118 @@ +# Combine the policy loss with the decoupled draft (MTP) loss. +# +# Mirrors NeMo-RL's ``DraftLossWrapper`` +# (https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/wrapper.py): +# the policy loss and the draft loss are computed independently and combined as +# ``policy_loss + w * draft_loss``. Because the draft head is fed *detached* +# trunk hidden states and a *detached* teacher distribution, the draft term only +# produces gradients for the draft-head parameters — the addition is just a +# convenient single backward, not an entanglement of the two objectives. + +from dataclasses import dataclass +from typing import Optional + +import torch + +from skyrl.backends.skyrl_train.mtp.soft_ce import ( + build_teacher_logits, + draft_hard_ce, + draft_soft_ce, + shift_mask_for_mtp, +) + + +@dataclass +class DraftLossConfig: + """Resolved knobs for the draft loss (populated from MegatronConfig).""" + + loss_weight: float = 0.1 + loss_type: str = "soft_ce" # "soft_ce" | "hard_ce" + + +def combine_policy_and_draft_loss( + policy_loss: torch.Tensor, + student_logits_per_layer: list[torch.Tensor], + main_logits: torch.Tensor, + mask: torch.Tensor, + cfg: DraftLossConfig, + hard_labels: Optional[torch.Tensor] = None, + global_normalization_factor: Optional[torch.Tensor] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + vocab_start_index: Optional[int] = None, +) -> tuple[torch.Tensor, dict]: + """Return ``policy_loss + w * draft_loss`` and a metrics dict. + + Args: + policy_loss: The already-computed scalar policy loss. + student_logits_per_layer: One ``[batch, seq, vocab(/tp)]`` draft-logits + tensor per MTP depth, aligned (same de-padding) with ``main_logits``. + main_logits: ``[batch, seq, vocab(/tp)]`` policy logits — the soft-CE + teacher source. + mask: ``[batch, seq]`` token mask over the region the draft is trained on. + cfg: Draft loss configuration. + hard_labels: ``[batch, seq]`` base next-token labels (``seq[t+1]`` at ``t``). + Rolled internally per MTP depth. Required when ``cfg.loss_type == "hard_ce"``. + global_normalization_factor: Optional global valid-token count for the + masked-mean denominator (cross-microbatch / DP correct reduction). + vocab_parallel_group: TP group when logits are vocab-sharded. + vocab_start_index: This rank's vocab offset (required for vocab-parallel + hard CE). + + Returns: + ``(combined_loss, metrics)`` where ``metrics`` carries the (detached) + draft loss for logging. + """ + draft_losses = [] + for layer_idx, student_logits in enumerate(student_logits_per_layer): + layer_mask = shift_mask_for_mtp(mask, layer_idx) + if cfg.loss_type == "hard_ce": + assert hard_labels is not None, "hard_ce requires hard_labels" + layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) + draft_losses.append( + draft_hard_ce( + student_logits, + layer_labels, + layer_mask, + global_normalization_factor=global_normalization_factor, + vocab_parallel_group=vocab_parallel_group, + vocab_start_index=vocab_start_index, + ) + ) + else: + teacher_logits = build_teacher_logits(main_logits, layer_idx, detach=True) + draft_losses.append( + draft_soft_ce( + student_logits, + teacher_logits, + layer_mask, + global_normalization_factor=global_normalization_factor, + vocab_parallel_group=vocab_parallel_group, + ) + ) + + draft_loss = torch.stack(draft_losses).mean() + combined = policy_loss + cfg.loss_weight * draft_loss + metrics = { + "draft_loss": draft_loss.detach().item(), + "mtp_loss": draft_loss.detach().item(), # alias kept for existing dashboards + } + return combined, metrics + + +class DraftLossWrapper: + """Object form of :func:`combine_policy_and_draft_loss` for backends that + prefer a stateful seam (matches NeMo-RL's ``DraftLossWrapper`` shape). The + FSDP backend will reuse this unchanged once its adapter lands.""" + + def __init__(self, cfg: DraftLossConfig): + self.cfg = cfg + + def __call__(self, policy_loss, student_logits_per_layer, main_logits, mask, **kwargs): + return combine_policy_and_draft_loss( + policy_loss, + student_logits_per_layer, + main_logits, + mask, + self.cfg, + **kwargs, + ) diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py new file mode 100644 index 0000000000..d22dbfa013 --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -0,0 +1,116 @@ +# Decoupled capture of the trunk hidden states that feed the MTP/draft head. +# +# This is the SkyRL analogue of NeMo-RL's ``HiddenStateCapture`` +# (https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/draft/hidden_capture.py). +# +# Megatron's ``GPTModel`` runs its native MTP block (``self.mtp``) inside the forward whenever +# ``self.mtp_process`` is set. Crucially, ``MultiTokenPredictionBlock.forward`` passes the trunk +# hidden states through as the *first* chunk of its output, and ``process_mtp_loss`` returns that +# chunk as the hidden states for the *main* logits. So we must NOT tamper with the in-forward call: +# the main policy logits have to stay connected to the trunk for the policy loss to train it. +# +# Instead we: +# 1. register a forward pre-hook on ``self.mtp`` that simply *records* the keyword arguments +# Megatron built for it (input_ids, position_ids, the live trunk hidden states, rotary +# embeddings, attention mask, packed-seq params, the shared embedding, ...). The pre-hook does +# not modify the call, so the in-forward MTP run proceeds normally and the main logits stay +# coupled to the trunk. +# 2. after the forward, re-invoke ``self.mtp`` with the recorded arguments but with the trunk +# hidden states *detached* (when ``detach_trunk``). This second pass is the decoupled draft +# forward: its gradient reaches the MTP-head (and shared embedding/output) parameters but never +# the policy backbone. Re-using the recorded kwargs means we never have to reconstruct rotary +# embeddings or attention masks ourselves. +# +# The in-forward MTP layers are recomputed (their output is discarded by ``process_mtp_loss`` when +# no labels are passed), so there is a small extra-compute cost — one MTP block forward, typically a +# single transformer layer. This is a deliberate trade for robustness over reaching into GPTModel's +# rotary/mask construction. + +from __future__ import annotations + +from contextlib import contextmanager +from typing import List, Optional + + +def _unwrap_gpt_model(model): + """Return the underlying GPTModel from a (possibly DDP/Float16) wrapper.""" + from megatron.core.utils import unwrap_model + + return unwrap_model(model) + + +class MTPHiddenCapture: + """Record the MTP block's inputs during the forward, then replay it decoupled. + + Use as a context manager around the policy ``model(...)`` call. After the forward, + :meth:`compute_student_hidden_states` returns one hidden-state tensor per MTP depth + (``[seq, batch, hidden]`` in Megatron's internal layout), ready to be projected by the shared + output layer. + """ + + def __init__(self, model, detach_trunk: bool = True): + self.gpt_model = _unwrap_gpt_model(model) + self.detach_trunk = detach_trunk + self._args = None + self._kwargs = None + self._handles: list = [] + + @property + def mtp_num_layers(self) -> int: + return int(getattr(self.gpt_model.config, "mtp_num_layers", 0) or 0) + + def _pre_hook(self, _module, args, kwargs): + # Record (do not modify) the arguments Megatron built for the MTP block. + self._args = args + self._kwargs = dict(kwargs) + return None + + @contextmanager + def capture(self): + mtp = getattr(self.gpt_model, "mtp", None) + if mtp is None: + # Model has no MTP heads on this rank/stage; nothing to capture. + yield self + return + self._args = None + self._kwargs = None + self._handles.append(mtp.register_forward_pre_hook(self._pre_hook, with_kwargs=True)) + try: + yield self + finally: + for h in self._handles: + h.remove() + self._handles.clear() + + def compute_student_hidden_states(self) -> Optional[List]: + """Replay the MTP block on detached trunk hidden states and split per depth. + + Returns ``None`` if the block was never called (e.g. a non-post-process pipeline stage). + """ + if self._kwargs is None: + return None + import torch + + kwargs = dict(self._kwargs) + hidden = kwargs.get("hidden_states") + if hidden is not None and self.detach_trunk: + kwargs["hidden_states"] = hidden.detach() + + # MultiTokenPredictionBlock returns the concatenated hidden states + # [trunk; mtp_0; mtp_1; ...] along the sequence (dim 0). + captured = self.gpt_model.mtp(*self._args, **kwargs) + num_layers = self.mtp_num_layers + chunks = list(torch.chunk(captured, 1 + num_layers, dim=0)) + # Drop the (passthrough) trunk chunk; keep the per-MTP-depth chunks. + return chunks[1:] + + +@contextmanager +def maybe_capture_mtp_hidden(model, enabled: bool, detach_trunk: bool = True): + """Context manager returning an ``MTPHiddenCapture`` when ``enabled``, else ``None``.""" + if not enabled: + yield None + return + capture = MTPHiddenCapture(model, detach_trunk=detach_trunk) + with capture.capture(): + yield capture diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py new file mode 100644 index 0000000000..05896d55c2 --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -0,0 +1,207 @@ +# Explicit draft-head losses for decoupled Multi-Token Prediction (MTP). +# +# Ported from NeMo-RL's ``DraftCrossEntropyLossFn`` and the ``DRAFT`` branch of +# ``prepare_loss_input``: +# https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/loss_functions.py +# https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/utils.py +# +# The soft cross-entropy matches the forward-KL student gradient +# (``softmax(student) - softmax(teacher)``) and is the default supervision for +# the draft head: the teacher is the *policy's own* next-token distribution +# (detached), so training the draft head never pulls on the policy trunk. + +from typing import Optional + +import torch +import torch.distributed +import torch.nn.functional as F + +from skyrl.backends.skyrl_train.utils.torch_utils import masked_mean + + +def build_teacher_logits( + main_logits: torch.Tensor, + mtp_layer_number: int = 0, + detach: bool = True, +) -> torch.Tensor: + """Build the soft-distillation teacher for a given MTP depth. + + The policy's main logits at position ``t`` are the model's distribution over + ``seq[t + 1]``. MTP layer ``k`` (0-indexed) predicts ``seq[t + k + 2]`` at + position ``t``, whose teacher distribution (conditioned on the verified + prefix) is the main model's distribution at position ``t + k + 1`` — i.e. + ``main_logits`` rolled left by ``k + 1``. This generalizes NeMo-RL's + single-step roll (its Eagle draft is one layer, rolled by ``-1``). + + Args: + main_logits: ``[batch, seq, vocab]`` policy logits. + mtp_layer_number: 0-indexed MTP depth ``k``. + detach: Detach the teacher so no gradient flows back into the policy + trunk (the decoupling that makes this "draft" training). + + Returns: + ``[batch, seq, vocab]`` teacher logits, rolled and (optionally) detached. + The wrapped-around tail positions are invalid; the caller's loss mask + must zero them out (see ``shift_mask_for_mtp``). + """ + teacher = main_logits.detach() if detach else main_logits + return torch.roll(teacher, shifts=-(mtp_layer_number + 1), dims=1) + + +def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.Tensor: + """Roll a ``[batch, seq]`` loss mask to align with an MTP teacher/label. + + Positions that wrap around the sequence end have no valid target, so they + are zeroed. This mirrors how Megatron's ``roll_tensor`` zeros the wrapped + boundary inside ``process_mtp_loss``. + """ + shift = mtp_layer_number + 1 + rolled = torch.roll(mask, shifts=-shift, dims=1) + rolled[:, -shift:] = 0 + return rolled + + +class _VocabParallelSoftCrossEntropy(torch.autograd.Function): + """Soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` + when ``student``/``teacher`` logits are sharded across the vocab (TP) dim. + + Forward reduces across the tensor-parallel group so the softmax denominators + and the teacher/student dot product are computed over the *global* vocab. + Backward returns the classic cross-entropy gradient + ``softmax(student) - softmax(teacher)`` (the teacher is treated as a + constant; no gradient flows to it). This is the TP analogue of NeMo-RL's + ``DistributedCrossEntropy``. + """ + + @staticmethod + def forward(ctx, student_vp_logits, teacher_vp_logits, tp_group): + student = student_vp_logits.float() + teacher = teacher_vp_logits.float() + + def global_softmax(logits): + local_max = logits.max(dim=-1, keepdim=True).values + torch.distributed.all_reduce(local_max, op=torch.distributed.ReduceOp.MAX, group=tp_group) + shifted = logits - local_max + exp = shifted.exp() + sum_exp = exp.sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(sum_exp, group=tp_group) + log_sum_exp = local_max + sum_exp.log() + return exp / sum_exp, log_sum_exp + + student_softmax, student_lse = global_softmax(student) + teacher_softmax, _ = global_softmax(teacher) + + # dot = sum_v teacher_softmax_v * student_logit_v (global over vocab) + dot = (teacher_softmax * student).sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(dot, group=tp_group) + + # soft CE = sum_v teacher_p_v * (log_sum_exp - student_logit_v) + # = log_sum_exp - dot (since sum_v teacher_p_v = 1) + per_token_loss = (student_lse - dot).squeeze(-1) + + ctx.save_for_backward(student_softmax, teacher_softmax) + return per_token_loss + + @staticmethod + def backward(ctx, grad_output): + student_softmax, teacher_softmax = ctx.saved_tensors + grad_student = (student_softmax - teacher_softmax) * grad_output.unsqueeze(-1) + return grad_student.to(student_softmax.dtype), None, None + + +def draft_soft_ce( + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + mask: torch.Tensor, + global_normalization_factor: Optional[torch.Tensor] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, +) -> torch.Tensor: + """Masked-mean soft cross-entropy between draft (student) and policy (teacher). + + Args: + student_logits: ``[batch, seq, vocab(/tp)]`` draft-head logits (require grad). + teacher_logits: ``[batch, seq, vocab(/tp)]`` policy logits (detached). + mask: ``[batch, seq]`` token mask (1 for supervised tokens). + global_normalization_factor: Optional global valid-token count used as the + masked-mean denominator (for correct cross-microbatch / DP reduction). + When ``None``, uses the local masked mean. + vocab_parallel_group: TP group when logits are vocab-sharded; ``None`` for + full-vocab logits (single device / FSDP). + + Returns: + Scalar loss. + """ + if vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1: + per_token_loss = _VocabParallelSoftCrossEntropy.apply( + student_logits, teacher_logits, vocab_parallel_group + ) + else: + teacher_probs = F.softmax(teacher_logits.float(), dim=-1) + student_log_probs = F.log_softmax(student_logits.float(), dim=-1) + per_token_loss = -(teacher_probs * student_log_probs).sum(dim=-1) + + if global_normalization_factor is not None: + return (per_token_loss * mask).sum() / global_normalization_factor.clamp(min=1.0) + return masked_mean(per_token_loss, mask) + + +def _onehot_vp_logits( + labels: torch.Tensor, + like: torch.Tensor, + vocab_start_index: int, +) -> torch.Tensor: + """Build vocab-parallel logits whose global softmax is a one-hot over ``labels``. + + Each rank holds ``vocab_size`` columns starting at ``vocab_start_index``. The + column matching the label (on whichever rank owns it) is set high and all + others low, so a *global* softmax across the TP group recovers the one-hot + distribution. Reused to express hard cross-entropy as soft cross-entropy with + a one-hot teacher. + """ + vocab_size = like.shape[-1] + local_idx = labels.long() - vocab_start_index # [batch, seq] + holds = (local_idx >= 0) & (local_idx < vocab_size) + onehot = torch.full_like(like, -30.0) + safe_idx = local_idx.clamp(0, vocab_size - 1).unsqueeze(-1) + hot = torch.where(holds.unsqueeze(-1), 30.0, -30.0).to(like.dtype) + onehot.scatter_(-1, safe_idx, hot) + return onehot + + +def draft_hard_ce( + student_logits: torch.Tensor, + labels: torch.Tensor, + mask: torch.Tensor, + global_normalization_factor: Optional[torch.Tensor] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + vocab_start_index: Optional[int] = None, +) -> torch.Tensor: + """Masked-mean hard cross-entropy of the draft head against next-token labels. + + This is the standard MTP-pretraining objective (supervise against the actual + future token rather than the policy distribution). Use when + ``mtp_loss_type="hard_ce"``. + + Args: + student_logits: ``[batch, seq, vocab(/tp)]`` draft-head logits. + labels: ``[batch, seq]`` target token ids (already rolled for this MTP depth). + mask: ``[batch, seq]`` token mask. + global_normalization_factor: see :func:`draft_soft_ce`. + vocab_parallel_group: TP group when logits are vocab-sharded. + vocab_start_index: This rank's vocab offset (required when TP-sharded). + """ + if vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1: + assert vocab_start_index is not None, "vocab_start_index is required for vocab-parallel hard CE" + # Hard CE == soft CE with a one-hot teacher; reuse the distributed soft-CE + # path so the global (TP) softmax / gradient logic lives in one place. + teacher_onehot = _onehot_vp_logits(labels, student_logits, vocab_start_index) + per_token_loss = _VocabParallelSoftCrossEntropy.apply( + student_logits, teacher_onehot, vocab_parallel_group + ) + else: + log_probs = F.log_softmax(student_logits.float(), dim=-1) + per_token_loss = -log_probs.gather(dim=-1, index=labels.unsqueeze(-1).long()).squeeze(-1) + + if global_normalization_factor is not None: + return (per_token_loss * mask).sum() / global_normalization_factor.clamp(min=1.0) + return masked_mean(per_token_loss, mask) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index e70836743b..a249ee137b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -29,42 +29,21 @@ setup_per_microbatch_replay_backward, setup_per_microbatch_replay_forward, ) +from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits +from skyrl.backends.skyrl_train.mtp.hidden_capture import maybe_capture_mtp_hidden +from skyrl.backends.skyrl_train.mtp.soft_ce import ( + build_teacher_logits, + draft_hard_ce, + draft_soft_ce, + shift_mask_for_mtp, +) from skyrl.backends.skyrl_train.utils.torch_utils import ( - build_mtp_loss_mask, build_mtp_next_token_labels, masked_mean, ) from skyrl.train.config import TrainerConfig -def _mtp_logits_passthrough( - *, - hidden_states, - output_layer, - output_weight, - runtime_gather_output=None, - scale_logits=None, - **kwargs, -): - """``output_processor`` hook for GPTModel that returns logits instead of the LM loss. - - When Multi-Token Prediction is active we pass ``labels``/``loss_mask`` into ``GPTModel.forward`` - so that ``process_mtp_loss`` runs and injects the MTP loss into the autograd graph (via - ``MTPLossAutoScaler`` on ``hidden_states``). But passing labels would normally make the model - return the *main* LM loss, whereas SkyRL needs the *logits* to compute its own PPO/GRPO/CE loss. - - GPTModel invokes this hook after ``process_mtp_loss`` and returns whatever it produces, short- - circuiting the labels->loss path. We reproduce the stock ``labels is None`` logits path exactly - (see ``GPTModel._postprocess``) so all downstream SkyRL logic is unchanged: vocab-parallel - logits in ``[batch, seq, vocab]`` layout. - """ - logits, _ = output_layer(hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output) - if scale_logits is not None: - logits = scale_logits(logits) - # [s b h] => [b s h], matching GPTModel's non-loss return. - return logits.transpose(0, 1).contiguous() - - class MegatronModelWrapper: def __init__( self, @@ -259,30 +238,22 @@ def forward_backward_mini_batch( forward_backward_func = get_forward_backward_func() # ------------------------------------------------------------------ - # Multi-Token Prediction (MTP) + # Multi-Token Prediction (MTP) — decoupled draft-head training. # ------------------------------------------------------------------ - # If the model was built with MTP heads (policy.megatron_config.mtp_num_layers), train them - # alongside the policy by feeding next-token labels into GPTModel.forward so Megatron's - # process_mtp_loss runs and injects the MTP gradient via MTPLossAutoScaler. We only do this - # when actually training (not forward_only / logprob-only passes). + # If the model was built with native MTP heads (policy.megatron_config.mtp_num_layers), + # train them with an explicit, decoupled loss instead of Megatron's in-forward + # process_mtp_loss / MTPLossAutoScaler path. We keep the heads running inside GPTModel.forward + # (so we don't have to reconstruct rotary embeddings) but pass NO labels, so process_mtp_loss + # short-circuits and no MTP gradient is coupled onto the trunk. A forward hook captures the + # MTP block's hidden states (with its trunk input optionally detached) and we project + score + # them ourselves. Only active during training (not forward_only / logprob-only passes). model_config = get_model_config(self.actor_module[0]) mtp_enabled = (not forward_only) and bool(getattr(model_config, "mtp_num_layers", None)) - if mtp_enabled: - from megatron.core.transformer.multi_token_prediction import ( - MTPLossAutoScaler, - MTPLossLoggingHelper, - ) - - # Match the pipeline schedule's per-microbatch loss division (Megatron divides the - # main loss by num_microbatches before backward). The MTP autoscaler injects a constant - # gradient independent of the main loss value, so we scale it the same way to keep the - # MTP aux-loss gradient commensurate with the policy-loss gradient. mtp_loss_scaling_factor - # (set on the provider) is the higher-level knob for the MTP/policy loss ratio. - MTPLossAutoScaler.set_loss_scale( - torch.tensor(1.0 / max(1, len(micro_batches)), device=torch.cuda.current_device()) - ) - if "values" in MTPLossLoggingHelper.tracker: - MTPLossLoggingHelper.clean_loss_in_tracker() + mcfg = self.cfg.policy.megatron_config + mtp_loss_type = getattr(mcfg, "mtp_loss_type", "soft_ce") + mtp_loss_weight = float(getattr(mcfg, "mtp_loss_weight", 0.1)) + mtp_detach_trunk = bool(getattr(mcfg, "mtp_detach_trunk", True)) + mtp_detach_shared_output = bool(getattr(mcfg, "mtp_detach_shared_output", False)) # Resolve loss function resolved_loss_name = loss_fn if loss_fn is not None else self.cfg.algorithm.policy_loss_type @@ -341,9 +312,54 @@ def loss_func(logits, data): rollout_logprobs=rollout_action_logprobs, ) + # --- Decoupled MTP / draft loss ------------------------------------------------- + # Score the (detached-input) MTP head against the policy's own next-token + # distribution (soft CE) or the ground-truth future tokens (hard CE), over every real + # token. Returns a local masked-mean scalar that is folded into the loss below with the + # same reduction treatment as the KL/entropy aux terms. + draft_loss = None + student_logits_list = data.get("mtp_student_logits") + if mtp_enabled and student_logits_list: + draft_mask = data["attention_mask"].to(logits.dtype) # [batch, seq_len] + vocab_size_tp = logits.shape[-1] + # Undo the in-place temperature scaling so the teacher is the true policy + # distribution (student logits are produced unscaled). + teacher_src = logits if temperature == 1.0 else logits * temperature + hard_labels = build_mtp_next_token_labels(sequences) if mtp_loss_type == "hard_ce" else None + + per_layer_losses = [] + for layer_idx, student_logits in enumerate(student_logits_list): + layer_mask = shift_mask_for_mtp(draft_mask, layer_idx) + if mtp_loss_type == "hard_ce": + layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) + per_layer_losses.append( + draft_hard_ce( + student_logits, + layer_labels, + layer_mask, + vocab_parallel_group=tp_grp, + vocab_start_index=tp_rank * vocab_size_tp, + ) + ) + else: + teacher_logits = build_teacher_logits(teacher_src, layer_idx, detach=True) + per_layer_losses.append( + draft_soft_ce( + student_logits, + teacher_logits, + layer_mask, + vocab_parallel_group=tp_grp, + ) + ) + draft_loss = torch.stack(per_layer_losses).mean() + # SFT path: cross_entropy loss (negative log likelihood) if resolved_loss_name == "cross_entropy": loss = policy_loss + if draft_loss is not None: + # Both terms are per-token means here; Megatron's /num_microbatches and the DDP + # DP+CP averaging reduce them consistently. + loss = loss + mtp_loss_weight * draft_loss # Compute elementwise loss for Tinker API (per-token NLL) with torch.no_grad(): @@ -386,6 +402,8 @@ def loss_func(logits, data): "response_length": num_actions, "loss_fn_outputs": loss_fn_outputs, } + if draft_loss is not None: + metrics["mtp_loss"] = draft_loss.detach().item() return loss, metrics # RL path: add optional KL/entropy terms @@ -428,6 +446,12 @@ def loss_func(logits, data): # Multiply by cp_size to counteract the unwanted CP averaging. cp_size = mpu.get_context_parallel_world_size() loss = policy_loss * grad_sum_correction_factor + (kl_loss_term - entropy_loss_term) * cp_size + # The decoupled MTP/draft loss is a per-token mean (like KL/entropy), so fold it in with + # the same cp_size correction. Its gradient only reaches the MTP-head parameters (and the + # shared output/embedding unless mtp_detach_shared_output) because both the trunk hidden + # states and the teacher distribution are detached. + if draft_loss is not None: + loss = loss + mtp_loss_weight * draft_loss * cp_size unscaled_loss = loss / grad_sum_correction_factor # Build per-sequence loss_fn_outputs with logprobs. @@ -457,6 +481,8 @@ def loss_func(logits, data): "policy_kl": kl_loss.detach().item(), "loss_fn_outputs": loss_fn_outputs, } + if draft_loss is not None: + metrics["mtp_loss"] = draft_loss.detach().item() for k, v in loss_metrics.items(): metrics["loss_metrics/" + k] = v return loss, metrics @@ -498,54 +524,54 @@ def forward_step(batch_iter, model): ) packed_seq_params = None - # When MTP is active, build next-token labels + loss mask in the SAME de-padded/packed - # space as new_sequences and pass them (plus the logits passthrough) so process_mtp_loss - # runs. We materialize labels on all ranks (pre_process=True) because they are consumed - # on the post_process / last pipeline stage, regardless of which stage owns the input. - mtp_kwargs = {} - if mtp_enabled: - labels_full = build_mtp_next_token_labels(sequences) - loss_mask_full = build_mtp_loss_mask(attention_mask) - if self.remove_microbatch_padding: - mtp_labels = preprocess_packed_seqs(labels_full, attention_mask, pre_process=True)[0] - mtp_loss_mask = preprocess_packed_seqs(loss_mask_full, attention_mask, pre_process=True)[0] - else: - mtp_labels = remove_left_padding(labels_full, attention_mask, position_ids, pre_process=True)[0] - mtp_loss_mask = remove_left_padding(loss_mask_full, attention_mask, position_ids, pre_process=True)[ - 0 - ] - mtp_kwargs = dict( - labels=mtp_labels, - loss_mask=mtp_loss_mask, - output_processor=_mtp_logits_passthrough, - ) + is_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=True) - outputs = model( - new_sequences, - new_position_ids, - new_attention_mask, - packed_seq_params=packed_seq_params, - **mtp_kwargs, - ) - - if self.remove_microbatch_padding: - outputs = postprocess_packed_seqs( - outputs, - packed_seq_params, + def depad(tensor): + """Recover [batch, seq_len, ...] padded layout from the internal layout, + matching exactly how the main logits are de-padded below.""" + if self.remove_microbatch_padding: + return postprocess_packed_seqs( + tensor, + packed_seq_params, + attention_mask, + micro_batch_size, + seq_len, + post_process=is_last_stage, + ) + return recover_left_padding( + tensor, + new_attention_mask, attention_mask, - micro_batch_size, seq_len, - post_process=mpu.is_pipeline_last_stage(ignore_virtual=True), + post_process=is_last_stage, ) - else: - outputs = recover_left_padding( - outputs, + + # Run the policy forward. When MTP is active, a pre-hook records the native MTP block's + # arguments (we pass NO labels, so GPTModel's process_mtp_loss short-circuits and the + # main logits stay coupled to the trunk; MTPLossAutoScaler never runs). + with maybe_capture_mtp_hidden(model, mtp_enabled, detach_trunk=mtp_detach_trunk) as capture: + outputs = model( + new_sequences, + new_position_ids, new_attention_mask, - attention_mask, - seq_len, - post_process=mpu.is_pipeline_last_stage(ignore_virtual=True), + packed_seq_params=packed_seq_params, ) + outputs = depad(outputs) + + # Replay the MTP block on *detached* trunk hidden states (decoupled draft forward), + # project through the shared output layer, and de-pad into the same + # [batch, seq_len, vocab/tp] layout as the main logits so the draft loss can be scored + # against the policy's own distribution (or hard labels) in loss_func. + if mtp_enabled and capture is not None: + student_hidden = capture.compute_student_hidden_states() + if student_hidden is not None: + gpt_model = capture.gpt_model + student_logits = project_mtp_hidden_to_logits( + student_hidden, gpt_model, detach_output_weight=mtp_detach_shared_output + ) + batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] + if rollout_expert_indices is not None: setup_per_microbatch_replay_backward() @@ -564,21 +590,8 @@ def forward_step(batch_iter, model): forward_only=forward_only, ) - # Surface the MTP loss for logging. process_mtp_loss accumulates a per-layer loss into - # MTPLossLoggingHelper.tracker (summed across microbatches) on the last pipeline stage. - # Attach the mean MTP loss to the first microbatch's metrics so it rides along the existing - # PP broadcast + DP reduction below. - if mtp_enabled and mpu.is_pipeline_last_stage(ignore_virtual=True): - from megatron.core.transformer.multi_token_prediction import ( - MTPLossLoggingHelper, - ) - - mtp_values = MTPLossLoggingHelper.tracker.get("values") - if mtp_values is not None and metrics_list and metrics_list[0] is not None: - mtp_mean = (mtp_values.sum() / (mtp_values.numel() * max(1, len(micro_batches)))).item() - metrics_list[0]["mtp_loss"] = mtp_mean - if mtp_values is not None: - MTPLossLoggingHelper.clean_loss_in_tracker() + # The decoupled MTP/draft loss is computed and logged per-microbatch inside loss_func + # (metric key "mtp_loss"); no MTPLossLoggingHelper plumbing is needed. # broadcast metrics to all pp ranks if not mpu.is_pipeline_last_stage(ignore_virtual=True): diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 384dbac6e7..ee6ad4ae7d 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -417,10 +417,16 @@ def init_configs( elif megatron_config.mtp_num_layers is not None: provider.mtp_num_layers = megatron_config.mtp_num_layers or None if getattr(provider, "mtp_num_layers", None): - provider.mtp_loss_scaling_factor = megatron_config.mtp_loss_scaling_factor + # The native MTP heads are still *built* (so the megatron-bridge weight mappings + # round-trip them to HF / vLLM), but SkyRL trains them with a decoupled, explicit + # loss (see MegatronModelWrapper) instead of Megatron's in-forward + # process_mtp_loss / MTPLossAutoScaler path. We therefore do not set + # provider.mtp_loss_scaling_factor; the explicit weight is megatron_config.mtp_loss_weight. logger.info( - f"MTP enabled: mtp_num_layers={provider.mtp_num_layers}, " - f"mtp_loss_scaling_factor={provider.mtp_loss_scaling_factor}" + f"MTP enabled (decoupled): mtp_num_layers={provider.mtp_num_layers}, " + f"mtp_loss_type={megatron_config.mtp_loss_type}, " + f"mtp_loss_weight={megatron_config.mtp_loss_weight}, " + f"mtp_detach_trunk={megatron_config.mtp_detach_trunk}" ) provider.finalize() diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index ea730df8a9..562de7c820 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -205,12 +205,35 @@ class MegatronConfig(BaseConfig): When MTP layers are present, SkyRL adds the Megatron MTP loss to the training step so the heads track the updated policy, and the heads are synced to vLLM for speculative decoding (set ``generator.inference_engine.speculative_config`` accordingly).""" - mtp_loss_scaling_factor: float = 0.1 - """Scaling coefficient for the MTP auxiliary loss (Megatron ``mtp_loss_scaling_factor``). - Only used when MTP layers are active. Higher values train the MTP heads more aggressively at the - cost of pulling on the shared trunk; the default matches Megatron/DeepSeek-V3.""" + mtp_loss_weight: float = 0.1 + """Weight ``w`` of the decoupled MTP/draft loss in ``policy_loss + w * draft_loss``. + + Replaces the role of Megatron's opaque ``mtp_loss_scaling_factor`` backward-scale. SkyRL now + trains the native MTP heads with an *explicit* loss on *detached* trunk hidden states (so the + draft gradient never pulls on the policy backbone) and adds it to the policy loss with this + weight. Only used when MTP layers are active.""" + mtp_loss_type: str = "soft_ce" + """Supervision for the MTP/draft head: + + - ``"soft_ce"`` (default): soft cross-entropy distillation against the policy's own detached, + rolled next-token distribution (matches NeMo-RL's draft training). + - ``"hard_ce"``: standard next-token cross-entropy against the ground-truth future tokens + (classic MTP-pretraining objective).""" + mtp_detach_trunk: bool = True + """If True (default, matches NeMo-RL), detach the trunk hidden states feeding the MTP head so + only the MTP-head parameters receive the draft gradient. If False, let the draft gradient flow + back into the policy backbone (closer to Megatron's native coupled MTP).""" + mtp_detach_shared_output: bool = False + """If True, also detach the shared embedding/output-layer weight from the draft gradient. + Default keeps the shared head trainable by the draft loss (only the trunk hidden is detached).""" + mtp_loss_scaling_factor: Optional[float] = None + """Deprecated alias for ``mtp_loss_weight`` (kept for back-compat). If set, it overrides + ``mtp_loss_weight``.""" def __post_init__(self): + # Back-compat: the old `mtp_loss_scaling_factor` knob now maps onto the explicit loss weight. + if self.mtp_loss_scaling_factor is not None: + self.mtp_loss_weight = self.mtp_loss_scaling_factor # Backfill defaults for any keys the user didn't override so an override dict # doesn't have to repeat every default just to set one value. if self.transformer_config_kwargs is None: diff --git a/tests/backends/skyrl_train/mtp/test_hidden_capture.py b/tests/backends/skyrl_train/mtp/test_hidden_capture.py new file mode 100644 index 0000000000..634ddff633 --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_hidden_capture.py @@ -0,0 +1,86 @@ +"""CPU unit tests for the MTP hidden-state capture / decoupled replay plumbing. + +These use a fake MTP block (no Megatron) to verify the capture records the block's +arguments and that the decoupled replay detaches the trunk hidden states so the +draft gradient never reaches the policy backbone. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_hidden_capture.py +""" + +import sys +import types + +import torch +import torch.nn as nn + +# Stub out megatron.core.utils.unwrap_model so hidden_capture imports on CPU. +_fake_mcore_utils = types.ModuleType("megatron.core.utils") +_fake_mcore_utils.unwrap_model = lambda m: m +sys.modules.setdefault("megatron", types.ModuleType("megatron")) +sys.modules.setdefault("megatron.core", types.ModuleType("megatron.core")) +sys.modules["megatron.core.utils"] = _fake_mcore_utils + +from skyrl.backends.skyrl_train.mtp.hidden_capture import MTPHiddenCapture # noqa: E402 + + +class _FakeMTPBlock(nn.Module): + """Mimics MultiTokenPredictionBlock: returns cat([trunk; mtp_0], dim=0).""" + + def __init__(self, hidden): + super().__init__() + self.w = nn.Parameter(torch.ones(hidden)) + + def forward(self, hidden_states, **kwargs): + mtp_0 = hidden_states * self.w # depends on params AND trunk hidden + return torch.cat([hidden_states, mtp_0], dim=0) + + +class _FakeGPT(nn.Module): + def __init__(self, hidden=4, mtp_num_layers=1): + super().__init__() + self.mtp = _FakeMTPBlock(hidden) + self.config = types.SimpleNamespace(mtp_num_layers=mtp_num_layers) + + +def _run(detach_trunk): + gpt = _FakeGPT() + capture = MTPHiddenCapture(gpt, detach_trunk=detach_trunk) + s, b, h = 3, 2, 4 + trunk = torch.randn(s, b, h, requires_grad=True) + with capture.capture(): + # Simulate GPTModel's in-forward MTP call (records kwargs via the pre-hook). + _ = gpt.mtp(hidden_states=trunk, position_ids=torch.zeros(b, s)) + student = capture.compute_student_hidden_states() + return gpt, trunk, student + + +def test_capture_returns_one_chunk_per_mtp_depth(): + _, _, student = _run(detach_trunk=True) + assert student is not None and len(student) == 1 + assert student[0].shape == (3, 2, 4) + + +def test_replay_detaches_trunk_when_detach_trunk_true(): + gpt, trunk, student = _run(detach_trunk=True) + student[0].sum().backward() + # The MTP-head parameter receives gradient... + assert gpt.mtp.w.grad is not None and gpt.mtp.w.grad.abs().sum() > 0 + # ...but the trunk hidden states do NOT (decoupled). + assert trunk.grad is None + + +def test_replay_couples_trunk_when_detach_trunk_false(): + gpt, trunk, student = _run(detach_trunk=False) + student[0].sum().backward() + assert gpt.mtp.w.grad is not None and gpt.mtp.w.grad.abs().sum() > 0 + # With detach disabled, the gradient flows back into the trunk hidden states. + assert trunk.grad is not None and trunk.grad.abs().sum() > 0 + + +def test_no_capture_when_block_absent(): + gpt = _FakeGPT() + gpt.mtp = None + capture = MTPHiddenCapture(gpt, detach_trunk=True) + with capture.capture(): + pass + assert capture.compute_student_hidden_states() is None diff --git a/tests/backends/skyrl_train/mtp/test_soft_ce.py b/tests/backends/skyrl_train/mtp/test_soft_ce.py new file mode 100644 index 0000000000..61babf2698 --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_soft_ce.py @@ -0,0 +1,116 @@ +"""CPU unit tests for the decoupled MTP draft losses and loss combination. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_soft_ce.py +""" + +import torch +import torch.nn.functional as F + +from skyrl.backends.skyrl_train.mtp.draft_loss_wrapper import ( + DraftLossConfig, + combine_policy_and_draft_loss, +) +from skyrl.backends.skyrl_train.mtp.soft_ce import ( + build_teacher_logits, + draft_hard_ce, + draft_soft_ce, + shift_mask_for_mtp, +) + + +def test_soft_ce_matches_reference(): + torch.manual_seed(0) + student = torch.randn(2, 5, 7, requires_grad=True) + teacher = torch.randn(2, 5, 7) + mask = torch.ones(2, 5) + + ref = -(F.softmax(teacher, -1) * F.log_softmax(student, -1)).sum(-1) + ref_mm = (ref * mask).sum() / mask.sum() + got = draft_soft_ce(student, teacher, mask) + assert torch.allclose(got, ref_mm, atol=1e-6) + + +def test_soft_ce_gradient_is_softmax_difference(): + # d/d student of soft CE is softmax(student) - softmax(teacher), spread over the mask mean. + torch.manual_seed(1) + student = torch.randn(2, 4, 6, requires_grad=True) + teacher = torch.randn(2, 4, 6) + mask = torch.ones(2, 4) + + draft_soft_ce(student, teacher, mask).backward() + n = mask.sum() + expected = (F.softmax(student.detach(), -1) - F.softmax(teacher, -1)) * (mask.unsqueeze(-1) / n) + assert torch.allclose(student.grad, expected, atol=1e-6) + + +def test_soft_ce_respects_mask(): + student = torch.randn(1, 3, 5, requires_grad=True) + teacher = torch.randn(1, 3, 5) + mask = torch.tensor([[1.0, 0.0, 1.0]]) + # Masked-out position must not affect the loss value. + teacher_alt = teacher.clone() + teacher_alt[0, 1] = torch.randn(5) + a = draft_soft_ce(student, teacher, mask) + b = draft_soft_ce(student, teacher_alt, mask) + assert torch.allclose(a, b, atol=1e-6) + + +def test_hard_ce_matches_reference(): + torch.manual_seed(2) + student = torch.randn(2, 5, 7, requires_grad=True) + labels = torch.randint(0, 7, (2, 5)) + mask = torch.ones(2, 5) + + got = draft_hard_ce(student, labels, mask) + ref = (-F.log_softmax(student, -1).gather(-1, labels.unsqueeze(-1)).squeeze(-1) * mask).sum() / mask.sum() + assert torch.allclose(got, ref, atol=1e-6) + + +def test_build_teacher_logits_rolls_and_detaches(): + ml = torch.arange(2 * 4 * 3, dtype=torch.float, requires_grad=True).reshape(2, 4, 3) + t0 = build_teacher_logits(ml, mtp_layer_number=0) + t1 = build_teacher_logits(ml, mtp_layer_number=1) + assert torch.equal(t0, torch.roll(ml.detach(), -1, dims=1)) + assert torch.equal(t1, torch.roll(ml.detach(), -2, dims=1)) + assert not t0.requires_grad + + +def test_shift_mask_zeros_boundary(): + m = torch.ones(1, 4) + assert shift_mask_for_mtp(m, 0).tolist() == [[1.0, 1.0, 1.0, 0.0]] + assert shift_mask_for_mtp(m, 1).tolist() == [[1.0, 1.0, 0.0, 0.0]] + + +def test_combine_is_policy_plus_weighted_draft(): + torch.manual_seed(3) + policy_loss = torch.tensor(2.0, requires_grad=True) + student = torch.randn(2, 5, 7, requires_grad=True) + main_logits = torch.randn(2, 5, 7) + mask = torch.ones(2, 5) + cfg = DraftLossConfig(loss_weight=0.5, loss_type="soft_ce") + + combined, metrics = combine_policy_and_draft_loss( + policy_loss, [student], main_logits, mask, cfg + ) + # Recompute the draft term independently and check the combination identity. + teacher = build_teacher_logits(main_logits, 0) + draft = draft_soft_ce(student, teacher, shift_mask_for_mtp(mask, 0)) + assert torch.allclose(combined, policy_loss + 0.5 * draft, atol=1e-6) + assert abs(metrics["mtp_loss"] - draft.item()) < 1e-6 + + +def test_combine_hard_ce_uses_rolled_labels(): + torch.manual_seed(4) + policy_loss = torch.tensor(1.0) + student = torch.randn(1, 5, 7, requires_grad=True) + main_logits = torch.randn(1, 5, 7) + mask = torch.ones(1, 5) + base_labels = torch.randint(0, 7, (1, 5)) + cfg = DraftLossConfig(loss_weight=1.0, loss_type="hard_ce") + + combined, _ = combine_policy_and_draft_loss( + policy_loss, [student], main_logits, mask, cfg, hard_labels=base_labels + ) + layer_labels = torch.roll(base_labels, shifts=-1, dims=1) + draft = draft_hard_ce(student, layer_labels, shift_mask_for_mtp(mask, 0)) + assert torch.allclose(combined, policy_loss + draft, atol=1e-6) diff --git a/tests/train/test_mtp_config.py b/tests/train/test_mtp_config.py index 47469c724d..ea20734e9d 100644 --- a/tests/train/test_mtp_config.py +++ b/tests/train/test_mtp_config.py @@ -11,13 +11,28 @@ def test_megatron_config_mtp_defaults(): cfg = MegatronConfig() # None => honor the model's own num_nextn_predict_layers (no SkyRL override). assert cfg.mtp_num_layers is None - assert cfg.mtp_loss_scaling_factor == 0.1 + # Decoupled draft-training defaults. + assert cfg.mtp_loss_weight == 0.1 + assert cfg.mtp_loss_type == "soft_ce" + assert cfg.mtp_detach_trunk is True + assert cfg.mtp_detach_shared_output is False + # Deprecated alias is unset by default. + assert cfg.mtp_loss_scaling_factor is None def test_megatron_config_mtp_overrides_parse(): - cfg = build_nested_dataclass(MegatronConfig, {"mtp_num_layers": 2, "mtp_loss_scaling_factor": 0.3}) + cfg = build_nested_dataclass( + MegatronConfig, {"mtp_num_layers": 2, "mtp_loss_weight": 0.3, "mtp_loss_type": "hard_ce"} + ) assert cfg.mtp_num_layers == 2 - assert cfg.mtp_loss_scaling_factor == 0.3 + assert cfg.mtp_loss_weight == 0.3 + assert cfg.mtp_loss_type == "hard_ce" + + +def test_megatron_config_mtp_loss_scaling_factor_back_compat(): + # The deprecated scaling-factor knob maps onto the explicit loss weight. + cfg = build_nested_dataclass(MegatronConfig, {"mtp_loss_scaling_factor": 0.25}) + assert cfg.mtp_loss_weight == 0.25 def test_megatron_config_mtp_force_disable(): From c25ddd5ed04464f196781340c020976934fb3dc8 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Fri, 5 Jun 2026 21:51:59 +0000 Subject: [PATCH 03/28] update megatron model wrapper, pass in position ids for MTP --- .../skyrl_train/mtp/hidden_capture.py | 10 +++++ .../megatron/megatron_model_wrapper.py | 39 ++++++++++++------- 2 files changed, 36 insertions(+), 13 deletions(-) diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index d22dbfa013..cde0043d14 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -54,6 +54,7 @@ def __init__(self, model, detach_trunk: bool = True): self._args = None self._kwargs = None self._handles: list = [] + self._prev_training = None @property def mtp_num_layers(self) -> int: @@ -74,6 +75,13 @@ def capture(self): return self._args = None self._kwargs = None + # Run the MTP block in eval mode for both the in-forward pass and the replay. Megatron's + # full-activation-recompute path (recompute_granularity='full' and module.training) routes + # through CheckpointFunction, which cannot save the non-tensor PackedSeqParams for backward. + # Eval skips that path (dropout is 0 in these configs; gradients still flow), and MTP is a + # single tiny layer so keeping its activations is negligible. + self._prev_training = mtp.training + mtp.eval() self._handles.append(mtp.register_forward_pre_hook(self._pre_hook, with_kwargs=True)) try: yield self @@ -81,6 +89,8 @@ def capture(self): for h in self._handles: h.remove() self._handles.clear() + if self._prev_training: + mtp.train() def compute_student_hidden_states(self) -> Optional[List]: """Replay the MTP block on detached trunk hidden states and split per depth. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index a249ee137b..682f39cef7 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -514,7 +514,18 @@ def forward_step(batch_iter, model): pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True), ) new_attention_mask = None - new_position_ids = None + # The trunk ignores position_ids for RoPE + THD packing (rotary comes from + # packed_seq_params), so SkyRL normally passes None. But the native MTP block rolls + # and re-embeds position_ids per depth, so when MTP is active we must supply them in + # the packed layout. This is harmless to the main logits. + if mtp_enabled: + new_position_ids = preprocess_packed_seqs( + position_ids, + attention_mask, + pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True), + )[0] + else: + new_position_ids = None else: new_sequences, new_attention_mask, new_position_ids = remove_left_padding( sequences, @@ -549,6 +560,7 @@ def depad(tensor): # Run the policy forward. When MTP is active, a pre-hook records the native MTP block's # arguments (we pass NO labels, so GPTModel's process_mtp_loss short-circuits and the # main logits stay coupled to the trunk; MTPLossAutoScaler never runs). + student_hidden = None with maybe_capture_mtp_hidden(model, mtp_enabled, detach_trunk=mtp_detach_trunk) as capture: outputs = model( new_sequences, @@ -556,21 +568,22 @@ def depad(tensor): new_attention_mask, packed_seq_params=packed_seq_params, ) + # Replay the MTP block on *detached* trunk hidden states (decoupled draft forward) + # while still inside the capture context (so the MTP block stays in eval mode). + if mtp_enabled and capture is not None: + student_hidden = capture.compute_student_hidden_states() + gpt_model = capture.gpt_model outputs = depad(outputs) - # Replay the MTP block on *detached* trunk hidden states (decoupled draft forward), - # project through the shared output layer, and de-pad into the same - # [batch, seq_len, vocab/tp] layout as the main logits so the draft loss can be scored - # against the policy's own distribution (or hard labels) in loss_func. - if mtp_enabled and capture is not None: - student_hidden = capture.compute_student_hidden_states() - if student_hidden is not None: - gpt_model = capture.gpt_model - student_logits = project_mtp_hidden_to_logits( - student_hidden, gpt_model, detach_output_weight=mtp_detach_shared_output - ) - batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] + # Project the decoupled MTP hidden states through the shared output layer and de-pad into + # the same [batch, seq_len, vocab/tp] layout as the main logits, so the draft loss can be + # scored against the policy's own distribution (or hard labels) in loss_func. + if student_hidden is not None: + student_logits = project_mtp_hidden_to_logits( + student_hidden, gpt_model, detach_output_weight=mtp_detach_shared_output + ) + batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] if rollout_expert_indices is not None: setup_per_microbatch_replay_backward() From 45931e6d72e93138c69a2be985ad21faaf0a0676 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Sat, 6 Jun 2026 20:17:11 +0000 Subject: [PATCH 04/28] add mtp spec-decode config, tests, spec decoding acceptance rate logging --- .../skyrl_train/inference_engines/base.py | 8 ++ .../inference_engine_client.py | 16 ++++ .../ray_wrapped_inference_engine.py | 3 + .../vllm/spec_decode_metrics.py | 79 ++++++++++++++++ .../inference_engines/vllm/vllm_engine.py | 50 ++++++++++- skyrl/train/config/__init__.py | 2 + skyrl/train/config/config.py | 26 ++++++ skyrl/train/trainer.py | 3 + .../mtp/test_spec_decode_metrics.py | 89 +++++++++++++++++++ tests/train/test_mtp_config.py | 40 ++++++++- 10 files changed, 313 insertions(+), 3 deletions(-) create mode 100644 skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py create mode 100644 tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py diff --git a/skyrl/backends/skyrl_train/inference_engines/base.py b/skyrl/backends/skyrl_train/inference_engines/base.py index 1ec593aae8..06d80882a1 100644 --- a/skyrl/backends/skyrl_train/inference_engines/base.py +++ b/skyrl/backends/skyrl_train/inference_engines/base.py @@ -174,6 +174,14 @@ async def resume_generation(self) -> None: """Resume generation after a pause, continuing any frozen in-flight requests.""" raise NotImplementedError + async def get_spec_decode_metrics(self) -> Optional[Dict[str, int]]: + """Return cumulative speculative-decoding counters, or None if unsupported. + + Non-abstract so backends without speculative decoding (or that can't expose the stats) work + unchanged. Keys when available: ``num_drafts``, ``num_draft_tokens``, ``num_accepted_tokens``. + """ + return None + @abstractmethod def tp_size(self) -> int: """Return the tensor parallel size of this inference engine.""" diff --git a/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py b/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py index eabd227a78..828fb32f24 100644 --- a/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py +++ b/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py @@ -383,6 +383,22 @@ async def resume_generation(self) -> None: """ await self._run_on_all_engines("resume_generation") + async def get_spec_decode_metrics(self) -> Optional[Dict[str, int]]: + """Sum cumulative speculative-decoding counters across all engines. + + Returns None if no engine reports spec-decode stats (e.g. speculative decoding disabled). + """ + per_engine = await self._run_on_all_engines("get_spec_decode_metrics") + totals: Dict[str, int] = {} + any_reported = False + for stats in per_engine: + if not stats: + continue + any_reported = True + for key, value in stats.items(): + totals[key] = totals.get(key, 0) + int(value) + return totals if any_reported else None + # ---------------------------- # HTTP endpoint related methods # ---------------------------- diff --git a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py index 00a92075e6..999a966174 100644 --- a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py +++ b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py @@ -82,6 +82,9 @@ async def pause_generation(self) -> None: async def resume_generation(self) -> None: return await self.inference_engine_actor.resume_generation.remote() + async def get_spec_decode_metrics(self): + return await self.inference_engine_actor.get_spec_decode_metrics.remote() + def create_ray_wrapped_inference_engines( num_inference_engines: int, diff --git a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py new file mode 100644 index 0000000000..637d56c1de --- /dev/null +++ b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py @@ -0,0 +1,79 @@ +# Cumulative speculative-decoding (MTP draft) acceptance accounting for vLLM v1. +# +# vLLM's async engine (the one SkyRL uses for generation + weight sync) does not expose +# ``get_metrics()``, so we attach a tiny custom stat logger that runs in the AsyncLLM frontend +# process (= the SkyRL engine actor) and accumulates the per-iteration spec-decode counters +# (``num_draft_tokens`` / ``num_accepted_tokens``). The engine reads these cumulative counters and +# the generator turns them into a per-step acceptance rate (accepted / drafted). + +from __future__ import annotations + +from typing import Optional + + +def make_spec_decode_stat_logger_class(): + """Return a ``StatLoggerBase`` subclass that sums spec-decode counts. + + Imported lazily and built as a class factory so this module stays importable without vLLM + (e.g. on CPU-only hosts / unit tests). + """ + from vllm.v1.metrics.loggers import StatLoggerBase + + class SpecDecodeStatLogger(StatLoggerBase): + """Accumulate cumulative spec-decode draft/accept counts across scheduler iterations.""" + + def __init__(self, vllm_config, engine_index: int = 0): + self.engine_index = engine_index + self.num_drafts = 0 + self.num_draft_tokens = 0 + self.num_accepted_tokens = 0 + + def record(self, scheduler_stats=None, iteration_stats=None, mm_cache_stats=None, engine_idx: int = 0): + stats = getattr(scheduler_stats, "spec_decoding_stats", None) if scheduler_stats is not None else None + if stats is not None: + self.num_drafts += stats.num_drafts + self.num_draft_tokens += stats.num_draft_tokens + self.num_accepted_tokens += stats.num_accepted_tokens + + def log_engine_initialized(self): + pass + + return SpecDecodeStatLogger + + +def sum_spec_decode_loggers(loggers) -> Optional[dict]: + """Sum cumulative counters across a list of ``SpecDecodeStatLogger`` instances.""" + if not loggers: + return None + return { + "num_drafts": sum(int(lg.num_drafts) for lg in loggers), + "num_draft_tokens": sum(int(lg.num_draft_tokens) for lg in loggers), + "num_accepted_tokens": sum(int(lg.num_accepted_tokens) for lg in loggers), + } + + +def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> tuple[dict, Optional[dict]]: + """Turn cumulative spec-decode counters into per-step (delta) metrics. + + Args: + cumulative: counters read this step (from the engines), or None if speculative decoding is + disabled / unsupported. + prev: the ``cumulative`` from the previous step (None on the first step). + + Returns: + ``(metrics, new_prev)`` where ``metrics`` has ``vllm/draft_*`` keys (empty when there are no + stats) and ``new_prev`` is the snapshot to pass back next step. + """ + if not cumulative: + return {}, prev + prev = prev or {} + drafted = cumulative.get("num_draft_tokens", 0) - prev.get("num_draft_tokens", 0) + accepted = cumulative.get("num_accepted_tokens", 0) - prev.get("num_accepted_tokens", 0) + metrics = { + "vllm/draft_num_draft_tokens": drafted, + "vllm/draft_num_accepted_tokens": accepted, + } + if drafted > 0: + # Acceptance rate = accepted draft tokens / total drafted tokens this step. + metrics["vllm/draft_acceptance_rate"] = accepted / drafted + return metrics, cumulative diff --git a/skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py b/skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py index 5f345037c0..35a6fc8008 100644 --- a/skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py +++ b/skyrl/backends/skyrl_train/inference_engines/vllm/vllm_engine.py @@ -111,12 +111,42 @@ def __init__(self, *args, bundle_indices: list = None, **kwargs): self._dp_size = kwargs.get("data_parallel_size", 1) self._is_lora = kwargs.get("enable_lora", False) + # Speculative-decoding (MTP draft) stat loggers, populated by the async engine. See + # spec_decode_metrics.py and get_spec_decode_metrics(). + self._spec_decode_loggers: list = [] + # Let subclass create the appropriate engine self.llm = self._create_engine(*args, **kwargs) # Weight loader is created by subclass after engine initialization self._weight_loader = None + async def get_spec_decode_metrics(self) -> Optional[Dict[str, int]]: + """Return cumulative spec-decode counters for this engine, or None if unavailable. + + Keys: ``num_drafts``, ``num_draft_tokens``, ``num_accepted_tokens``. The generator turns the + per-step delta of draft/accepted tokens into an acceptance rate. + """ + from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + sum_spec_decode_loggers, + ) + + if self._spec_decode_loggers: + return sum_spec_decode_loggers(self._spec_decode_loggers) + # Sync engine fallback: vllm.LLM exposes get_metrics(). + get_metrics = getattr(self.llm, "get_metrics", None) + if get_metrics is None: + return None + draft = accepted = None + for metric in get_metrics(): + if metric.name == "vllm:spec_decode_num_draft_tokens": + draft = int(metric.value) + elif metric.name == "vllm:spec_decode_num_accepted_tokens": + accepted = int(metric.value) + if draft is None and accepted is None: + return None + return {"num_drafts": 0, "num_draft_tokens": draft or 0, "num_accepted_tokens": accepted or 0} + def tp_size(self): return self._tp_size @@ -363,9 +393,25 @@ def _create_engine(self, *args, **kwargs): engine_args = vllm.AsyncEngineArgs(enable_log_requests=enable_log_requests, kv_cache_metrics=True, **kwargs) # Setup stat loggers for vLLM v1 if Ray Prometheus stats are enabled - stat_loggers = None + stat_loggers = [] if enable_ray_prometheus_stats: - stat_loggers = self._create_ray_prometheus_stat_loggers() + stat_loggers = list(self._create_ray_prometheus_stat_loggers()) + + # Always attach a lightweight spec-decode stat logger. It runs in this (frontend) process and + # accumulates draft/accept counts so the generator can report a per-step acceptance rate. + # No-op when speculative decoding is disabled (counts stay 0). + from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + make_spec_decode_stat_logger_class, + ) + + spec_logger_cls = make_spec_decode_stat_logger_class() + + def _spec_decode_logger_factory(vllm_config, engine_index: int = 0): + logger_instance = spec_logger_cls(vllm_config, engine_index) + self._spec_decode_loggers.append(logger_instance) + return logger_instance + + stat_loggers.append(_spec_decode_logger_factory) engine = vllm.AsyncLLMEngine.from_engine_args(engine_args, stat_loggers=stat_loggers) diff --git a/skyrl/train/config/__init__.py b/skyrl/train/config/__init__.py index fb81dec651..4f1b8c5b3e 100644 --- a/skyrl/train/config/__init__.py +++ b/skyrl/train/config/__init__.py @@ -21,6 +21,7 @@ MegatronTorchProfilerConfig, MixedPrecisionConfig, ModelConfig, + MTPConfig, OffPolicyCorrectionConfig, OptimizerConfig, PlacementConfig, @@ -54,6 +55,7 @@ "AlgorithmConfig", "GeneratorConfig", "InferenceEngineConfig", + "MTPConfig", "EnvironmentConfig", "ModelConfig", "SkyRLLoraConfig", diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 562de7c820..a542fe1acf 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -664,6 +664,31 @@ class EnvironmentConfig(BaseConfig): skyrl_gym: SkyRLGymConfig = field(default_factory=SkyRLGymConfig) +@dataclass +class MTPConfig(BaseConfig): + """High-level Multi-Token Prediction (MTP) control. + + This is the single user-facing knob for MTP. When ``enabled``, ``validate_cfg`` propagates it to + the training side (``policy.megatron_config.mtp_*`` — build and train ``num_speculative_tokens`` + native MTP heads with a decoupled draft loss) and the inference side + (``generator.inference_engine.speculative_config`` — vLLM MTP speculative decoding with the same + number of draft tokens). The trained heads are kept in sync with the policy via weight sync, so + the draft tracks the updated policy. + + The trunk hidden states and embeddings feeding the heads are always detached (decoupled) — that + is intrinsic to draft/Eagle training, not a tunable, so it is not exposed here.""" + + enabled: bool = False + """Whether to train MTP draft heads and use them for speculative decoding.""" + num_speculative_tokens: int = 1 + """Number of future tokens to predict = number of MTP heads / draft depth. Must be <= the number + of MTP layers the model ships (e.g. GLM-4.7-Flash / DeepSeek-V3 ``num_nextn_predict_layers``).""" + loss_type: str = "soft_ce" + """``"soft_ce"`` (distill against the policy's own next-token distribution) or ``"hard_ce"``.""" + loss_weight: float = 0.1 + """Weight ``w`` of the draft loss in ``policy_loss + w * draft_loss``.""" + + # --------------------------------------------------------------------------- # Trainer (top-level) # --------------------------------------------------------------------------- @@ -678,6 +703,7 @@ class TrainerConfig(BaseConfig): ref: RefConfig = field(default_factory=RefConfig) critic: CriticConfig = field(default_factory=CriticConfig) algorithm: AlgorithmConfig = field(default_factory=AlgorithmConfig) + mtp: MTPConfig = field(default_factory=MTPConfig) fully_async: FullyAsyncConfig = field(default_factory=FullyAsyncConfig) gradient_checkpointing: bool = True gradient_checkpointing_use_reentrant: bool = False diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 508b6b4635..b9ce0eee73 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -128,6 +128,9 @@ def __init__( self.all_metrics = {} self.all_timings = {} self.global_step = 0 + # Previous-step cumulative vLLM spec-decode (MTP draft) counters, used to report a per-step + # acceptance rate. See _record_spec_decode_metrics. + self._prev_spec_decode = None self._vllm_metrics_scraper: Optional[VLLMMetricsScraper] = ( VLLMMetricsScraper() if cfg.generator.inference_engine.enable_ray_prometheus_stats else None diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py new file mode 100644 index 0000000000..a69c3bf07c --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py @@ -0,0 +1,89 @@ +"""CPU unit tests for vLLM spec-decode (MTP draft) acceptance accounting. + +uv run --isolated --extra dev --extra megatron pytest tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py +""" + +import asyncio +import types + +from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + make_spec_decode_stat_logger_class, + sum_spec_decode_loggers, +) +from skyrl.train.generators.skyrl_gym_generator import SkyRLGymGenerator + + +def _sched(num_draft, num_accepted, num_drafts): + return types.SimpleNamespace( + spec_decoding_stats=types.SimpleNamespace( + num_drafts=num_drafts, num_draft_tokens=num_draft, num_accepted_tokens=num_accepted + ) + ) + + +def test_logger_accumulates_across_iterations(): + cls = make_spec_decode_stat_logger_class() + lg = cls(vllm_config=None, engine_index=0) + lg.record(scheduler_stats=_sched(10, 7, 5)) + lg.record(scheduler_stats=_sched(8, 6, 4)) + lg.record(scheduler_stats=types.SimpleNamespace(spec_decoding_stats=None)) # no spec this iter + assert (lg.num_draft_tokens, lg.num_accepted_tokens, lg.num_drafts) == (18, 13, 9) + + +def test_sum_spec_decode_loggers(): + cls = make_spec_decode_stat_logger_class() + a = cls(None, 0) + a.num_draft_tokens, a.num_accepted_tokens, a.num_drafts = 10, 6, 3 + b = cls(None, 1) + b.num_draft_tokens, b.num_accepted_tokens, b.num_drafts = 4, 2, 1 + assert sum_spec_decode_loggers([a, b]) == { + "num_drafts": 4, + "num_draft_tokens": 14, + "num_accepted_tokens": 8, + } + assert sum_spec_decode_loggers([]) is None + + +class _FakeClient: + def __init__(self, values): + self._values = list(values) + self._i = 0 + + async def get_spec_decode_metrics(self): + v = self._values[min(self._i, len(self._values) - 1)] + self._i += 1 + return v + + +def _make_generator(client): + gen = SkyRLGymGenerator.__new__(SkyRLGymGenerator) + gen.inference_engine_client = client + gen._prev_spec_decode = None + return gen + + +def test_generator_acceptance_rate_per_step_delta(): + # cumulative counters grow each step; the metric is the per-step delta ratio. + client = _FakeClient( + [ + {"num_drafts": 5, "num_draft_tokens": 10, "num_accepted_tokens": 6}, + {"num_drafts": 11, "num_draft_tokens": 22, "num_accepted_tokens": 15}, + ] + ) + gen = _make_generator(client) + + m1 = asyncio.run(gen._spec_decode_rollout_metrics()) + assert m1["vllm/draft_num_draft_tokens"] == 10 + assert m1["vllm/draft_num_accepted_tokens"] == 6 + assert abs(m1["vllm/draft_acceptance_rate"] - 0.6) < 1e-9 + + m2 = asyncio.run(gen._spec_decode_rollout_metrics()) + # deltas: drafted 22-10=12, accepted 15-6=9 -> 0.75 + assert m2["vllm/draft_num_draft_tokens"] == 12 + assert m2["vllm/draft_num_accepted_tokens"] == 9 + assert abs(m2["vllm/draft_acceptance_rate"] - 0.75) < 1e-9 + + +def test_generator_no_spec_decode_returns_empty(): + gen = _make_generator(_FakeClient([None])) + assert asyncio.run(gen._spec_decode_rollout_metrics()) == {} diff --git a/tests/train/test_mtp_config.py b/tests/train/test_mtp_config.py index ea20734e9d..1dcf46fb0a 100644 --- a/tests/train/test_mtp_config.py +++ b/tests/train/test_mtp_config.py @@ -3,8 +3,9 @@ uv run --isolated --extra dev pytest tests/train/test_mtp_config.py """ -from skyrl.train.config import InferenceEngineConfig, MegatronConfig +from skyrl.train.config import InferenceEngineConfig, MegatronConfig, MTPConfig, SkyRLTrainConfig from skyrl.train.config.config import build_nested_dataclass +from skyrl.train.utils.utils import _apply_mtp_config def test_megatron_config_mtp_defaults(): @@ -50,3 +51,40 @@ def test_inference_engine_speculative_config_parses_mtp_dict(): spec = {"method": "mtp", "num_speculative_tokens": 1} cfg = build_nested_dataclass(InferenceEngineConfig, {"speculative_config": spec}) assert cfg.speculative_config == spec + + +def test_mtp_config_defaults(): + cfg = MTPConfig() + assert cfg.enabled is False + assert cfg.num_speculative_tokens == 1 + assert cfg.loss_type == "soft_ce" + assert cfg.loss_weight == 0.1 + + +def test_apply_mtp_config_enabled_propagates_to_training_and_inference(): + cfg = SkyRLTrainConfig() + cfg.trainer.mtp.enabled = True + cfg.trainer.mtp.num_speculative_tokens = 2 + cfg.trainer.mtp.loss_weight = 0.25 + _apply_mtp_config(cfg) + assert cfg.trainer.policy.megatron_config.mtp_num_layers == 2 + assert cfg.trainer.policy.megatron_config.mtp_loss_weight == 0.25 + assert cfg.generator.inference_engine.speculative_config == { + "method": "mtp", + "num_speculative_tokens": 2, + } + + +def test_apply_mtp_config_disabled_force_disables_heads(): + cfg = SkyRLTrainConfig() + _apply_mtp_config(cfg) + assert cfg.trainer.policy.megatron_config.mtp_num_layers == 0 + assert cfg.generator.inference_engine.speculative_config is None + + +def test_apply_mtp_config_does_not_clobber_explicit_speculative_config(): + cfg = SkyRLTrainConfig() + cfg.trainer.mtp.enabled = True + cfg.generator.inference_engine.speculative_config = {"method": "mtp", "num_speculative_tokens": 5} + _apply_mtp_config(cfg) + assert cfg.generator.inference_engine.speculative_config["num_speculative_tokens"] == 5 From 15f757fc6831fa97fd943bfb5da08dc2b0fe6d83 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Sun, 7 Jun 2026 22:39:19 +0000 Subject: [PATCH 05/28] fix spec decode config not passing through to inference engine ray wrapper, add tests, qwen 3.5 script --- .../megatron/run_megatron_dapo_qwen3.5_2b.sh | 149 +++++++++++++++ ...run_megatron_dapo_qwen3.5_2b_specdecode.sh | 176 ++++++++++++++++++ .../ray_wrapped_inference_engine.py | 9 +- skyrl/backends/skyrl_train/mtp/adapter.py | 42 +++-- .../skyrl_train/mtp/hidden_capture.py | 56 ++++-- .../megatron/megatron_model_wrapper.py | 33 +++- .../workers/megatron/megatron_worker.py | 8 +- skyrl/backends/skyrl_train_backend.py | 1 + skyrl/train/entrypoints/main_base.py | 1 + skyrl/train/trainer.py | 22 +++ skyrl/train/utils/utils.py | 53 ++++++ .../test_inference_engine_client.py | 57 ++++++ .../skyrl_train/mtp/test_hidden_capture.py | 107 +++++++++-- .../mtp/test_spec_decode_metrics.py | 53 ++---- tests/train/test_mtp_config.py | 21 +++ 15 files changed, 693 insertions(+), 95 deletions(-) create mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh create mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh new file mode 100644 index 0000000000..e7a3128a91 --- /dev/null +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh @@ -0,0 +1,149 @@ +set -x + +# Colocated DAPO training+generation for Qwen3.5-2B (dense) on DAPO with Megatron. +# Runs on 1 node of 8xH100s (80GB each). +# +# NOTE: verify the exact HF repo id for the 2B model before running +# (e.g. `hf download Qwen/Qwen3.5-2B` / check https://huggingface.co/Qwen). +# +# Prepare data onto the fast local disk first: +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# Then launch: +# bash examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh + +MODEL_NAME="Qwen/Qwen3.5-2B" +# Use the fast, non-persistent local disk for data (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +# 2B is small: TP=1 inference per engine (no TP comm), one engine per GPU. +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 +LOGGER="wandb" # change to "console" to print to stdout + +CLIP_RATIO_LOW=0.2 +CLIP_RATIO_HIGH=0.28 +# use token mean loss reduction +LOSS_REDUCTION="token_mean" +# applies overlong filtering (but not soft overlong punishment) +APPLY_OVERLONG_FILTERING=true +# apply soft overlong punishment with custom trainer impl in main_dapo.py +OVERLONG_BUFFER_LEN=$((1024 * 4)) +OVERLONG_BUFFER_PENALTY_FACTOR=1.0 + +# other DAPO parameters +USE_KL_LOSS=false +TEMPERATURE=1.0 +TOP_P=1.0 +EVAL_TOP_P=0.7 +CLIP_RATIO_C=10.0 +MAX_PROMPT_LENGTH=$((1024 * 2)) +MAX_RESPONSE_LENGTH=$((1024 * 8)) + +# repro run parameters +# Quality-comparison profile: larger batch lowers reward-curve variance (~4x vs batch 32) +# so a spec-decode quality effect is detectable. rollout = 128 * 8 = 1024 seqs. +# MINI_BATCH_SIZE == TRAIN_BATCH_SIZE => 1 on-policy update/rollout. Keep IDENTICAL across no-spec/spec runs. +TRAIN_BATCH_SIZE=128 +MINI_BATCH_SIZE=32 +N_SAMPLES_PER_PROMPT=8 +EVAL_N_SAMPLES_PER_PROMPT=16 +ENFORCE_EAGER=true # cuda graphs can cause some instability +LR=1e-6 + +# megatron config -- Qwen3.5-2B is a dense model, so no expert parallelism. +# TP=2 (not 1): with 8K-token responses, TP=2 auto-enables sequence parallelism, +# which shards activations/vocab-logits across the TP group. TP=1 keeps full +# activations on each GPU and OOMs in the backward grad-sync at gpu_mem_util=0.5. +# TP=2, PP=1, CP=1 => DP=4. +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + + +# TIS parameters +TIS_IMP_RATIO_CAP=2.0 +TIS_TYPE=token + + +# Qwen3.5 flags +REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 +ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 +DISTRIBUTED_EXECUTOR_BACKEND="mp" +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="dual_clip" \ + trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ + trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + generator.sampling_params.temperature=$TEMPERATURE \ + generator.sampling_params.top_p=$TOP_P \ + generator.eval_sampling_params.top_p=$EVAL_TOP_P \ + generator.eval_sampling_params.temperature=$TEMPERATURE \ + generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=10 \ + trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ + trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=-1 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=$LR \ + trainer.policy.optimizer_config.num_warmup_steps=5 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.logger="$LOGGER" \ + trainer.project_name="qwen3_5_dapo_sd" \ + trainer.run_name="nosd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.export_path="/mnt/local_storage/exports/dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.hf_save_interval=300 \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_path="/mnt/local_storage/ckpts/dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + $@ diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh new file mode 100644 index 0000000000..a86dec2a7c --- /dev/null +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh @@ -0,0 +1,176 @@ +set -x + +# Colocated DAPO training+generation for Qwen3.5-2B (dense) on DAPO with Megatron, +# WITH Multi-Token Prediction (MTP) speculative decoding for faster rollout. +# +# This is the spec-decode counterpart of run_megatron_dapo_qwen3.5_2b.sh. Every +# knob below is IDENTICAL to the no-spec script (same batch sizes, LR, parallelism, +# sampling) so a reward-curve / throughput comparison is apples-to-apples. The ONLY +# difference is the `trainer.mtp.*` block at the bottom. +# +# What MTP on does (single high-level `trainer.mtp` knob, see skyrl/train/config/config.py): +# - Training side: builds + trains Qwen3.5-2B's native MTP head (the model ships +# `mtp_num_hidden_layers: 1`) with a decoupled draft loss (soft-CE distillation against +# the policy's own detached next-token distribution). The draft gradient never pulls on +# the policy backbone. +# - Inference side: enables vLLM MTP speculative decoding +# (`speculative_config={"method": "mtp", "num_speculative_tokens": 1}`). vLLM loads the +# MTP head from the same policy checkpoint, and SkyRL's weight sync keeps the draft head +# in sync with the trained policy each step. +# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. +# +# Runs on 1 node of 8xH100s (80GB each). +# +# Prepare data onto the fast local disk first: +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# Then launch: +# bash examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh + +MODEL_NAME="Qwen/Qwen3.5-2B" +# Use the fast, non-persistent local disk for data (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +# 2B is small: TP=1 inference per engine (no TP comm), one engine per GPU. +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 +LOGGER="wandb" # change to "console" to print to stdout + +CLIP_RATIO_LOW=0.2 +CLIP_RATIO_HIGH=0.28 +# use token mean loss reduction +LOSS_REDUCTION="token_mean" +# applies overlong filtering (but not soft overlong punishment) +APPLY_OVERLONG_FILTERING=true +# apply soft overlong punishment with custom trainer impl in main_dapo.py +OVERLONG_BUFFER_LEN=$((1024 * 4)) +OVERLONG_BUFFER_PENALTY_FACTOR=1.0 + +# other DAPO parameters +USE_KL_LOSS=false +TEMPERATURE=1.0 +TOP_P=1.0 +EVAL_TOP_P=0.7 +CLIP_RATIO_C=10.0 +MAX_PROMPT_LENGTH=$((1024 * 2)) +MAX_RESPONSE_LENGTH=$((1024 * 8)) + +# repro run parameters +# Quality-comparison profile: larger batch lowers reward-curve variance (~4x vs batch 32) +# so a spec-decode quality effect is detectable. rollout = 128 * 8 = 1024 seqs. +# MINI_BATCH_SIZE == TRAIN_BATCH_SIZE => 1 on-policy update/rollout. Keep IDENTICAL across no-spec/spec runs. +TRAIN_BATCH_SIZE=128 +MINI_BATCH_SIZE=32 +N_SAMPLES_PER_PROMPT=8 +EVAL_N_SAMPLES_PER_PROMPT=16 +ENFORCE_EAGER=true # cuda graphs can cause some instability +LR=1e-6 + +# megatron config -- Qwen3.5-2B is a dense model, so no expert parallelism. +# TP=2 (not 1): with 8K-token responses, TP=2 auto-enables sequence parallelism, +# which shards activations/vocab-logits across the TP group. TP=1 keeps full +# activations on each GPU and OOMs in the backward grad-sync at gpu_mem_util=0.5. +# TP=2, PP=1, CP=1 => DP=4. +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + + +# TIS parameters +TIS_IMP_RATIO_CAP=2.0 +TIS_TYPE=token + + +# Multi-Token Prediction (MTP) speculative decoding. +# Qwen3.5-2B ships 1 native MTP head (`mtp_num_hidden_layers: 1`), so draft depth = 1. +MTP_ENABLED=true +MTP_NUM_SPECULATIVE_TOKENS=1 +MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) +MTP_LOSS_WEIGHT=0.1 + + +# Qwen3.5 flags +REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 +ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 +DISTRIBUTED_EXECUTOR_BACKEND="mp" +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="dual_clip" \ + trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ + trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + generator.sampling_params.temperature=$TEMPERATURE \ + generator.sampling_params.top_p=$TOP_P \ + generator.eval_sampling_params.top_p=$EVAL_TOP_P \ + generator.eval_sampling_params.temperature=$TEMPERATURE \ + generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=10 \ + trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ + trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=-1 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=$LR \ + trainer.policy.optimizer_config.num_warmup_steps=5 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.mtp.enabled=$MTP_ENABLED \ + trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ + trainer.mtp.loss_type=$MTP_LOSS_TYPE \ + trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ + trainer.logger="$LOGGER" \ + trainer.project_name="qwen3_5_dapo_sd" \ + trainer.run_name="sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.export_path="/mnt/local_storage/exports/sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.hf_save_interval=300 \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_path="/mnt/local_storage/ckpts/sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + $@ diff --git a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py index 999a966174..48b693588c 100644 --- a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py +++ b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING, Any, Dict, List, Optional import ray from packaging import version @@ -119,6 +119,7 @@ def create_ray_wrapped_inference_engines( enable_return_routed_experts: bool = False, served_model_name: str | None = None, distributed_executor_backend: str = "ray", + speculative_config: Optional[Dict[str, Any]] = None, ) -> List[InferenceEngineInterface]: """ Create a list of RayWrappedInferenceEngine instances wrapping Ray actor handles to InferenceEngineInterface @@ -240,6 +241,12 @@ def create_ray_wrapped_inference_engines( if served_model_name is not None: other_kwargs["served_model_name"] = served_model_name + # Speculative decoding (e.g. MTP draft). Passed straight through to vLLM's + # AsyncEngineArgs.speculative_config so the colocated local engine path enables spec + # decode too (the standalone HTTP-server path handles this in build_vllm_cli_args). + if speculative_config is not None: + other_kwargs["speculative_config"] = speculative_config + # Launch one actor per DP rank for dp_rank in range(data_parallel_size): diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py index 78c2eb39ff..af33168a47 100644 --- a/skyrl/backends/skyrl_train/mtp/adapter.py +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -2,10 +2,19 @@ # # The loss (``soft_ce``), the combination (``draft_loss_wrapper``) and the # capture mechanism (``hidden_capture``) are backend-agnostic. Only two things -# differ per backend: (a) where the trunk hidden states are captured and (b) how -# the draft head turns those hidden states into vocab logits. This module pins -# down that contract. Megatron is implemented today; FSDP is a planned follow-up -# that will implement the same protocol and reuse everything else unchanged. +# differ per backend / draft style: (a) where the trunk hidden states are +# captured and (b) how the draft head turns those hidden states into vocab +# logits. ``DraftAdapter`` pins down that contract. +# +# Implemented today: native-MTP draft training on Megatron (capture the model's +# built-in ``self.mtp`` block and project through the shared output layer), which +# works for any model that ships native MTP heads regardless of base class +# (``GPTModel`` for DeepSeek/GLM/Qwen3-Next, ``MambaModel`` for Qwen3.5/NemotronH). +# +# Planned follow-ups implement the SAME protocol and reuse the loss/combination +# unchanged: (1) a NeMo-RL-style Eagle adapter that captures auxiliary policy +# hidden states and projects them through a separate Eagle draft model (for models +# WITHOUT native MTP heads); (2) an FSDP backend adapter. from __future__ import annotations @@ -14,18 +23,22 @@ @runtime_checkable class DraftAdapter(Protocol): - """Per-backend hooks for decoupled MTP training.""" + """Per-backend / per-draft-style hooks for decoupled MTP/draft training. + + An Eagle adapter would implement the same two methods: ``capture_context`` + grabs the auxiliary policy hidden states, and ``project_to_logits`` runs the + (separate) draft model's output head.""" def capture_context(self, model, detach_trunk: bool): ... def project_to_logits(self, hidden_states_per_layer, model) -> List: ... -def project_mtp_hidden_to_logits(hidden_states_per_layer, gpt_model, detach_output_weight: bool = False) -> List: - """Run the shared output layer on each captured MTP hidden-state chunk. +def project_mtp_hidden_to_logits(hidden_states_per_layer, model, detach_output_weight: bool = False) -> List: + """Run the model's shared output layer on each captured MTP hidden-state chunk. - Mirrors ``GPTModel._postprocess``: logits come out of the output layer in - ``[seq, batch, vocab/tp]`` and are transposed to ``[batch, seq, vocab/tp]`` + Mirrors the model's own ``_postprocess``: logits come out of the output layer + in ``[seq, batch, vocab/tp]`` and are transposed to ``[batch, seq, vocab/tp]`` to match the main-logits layout. The returned tensors are still in Megatron's *internal* (packed / left-removed) sequence layout; the caller applies the same de-padding transform used for the main logits so they align with the @@ -34,15 +47,16 @@ def project_mtp_hidden_to_logits(hidden_states_per_layer, gpt_model, detach_outp Args: hidden_states_per_layer: list of ``[seq, batch, hidden]`` tensors, one per MTP depth (from ``MTPHiddenCapture.compute_student_hidden_states``). - gpt_model: the unwrapped ``GPTModel`` (exposes ``output_layer`` and, when - embeddings are tied, ``shared_embedding_or_output_weight``). + model: the unwrapped Megatron model (``GPTModel``/``MambaModel``/...), + which exposes ``output_layer`` and, when embeddings are tied, + ``shared_embedding_or_output_weight``. Returns: list of ``[batch, seq, vocab/tp]`` student-logits tensors (internal layout). """ output_weight = None - if getattr(gpt_model, "share_embeddings_and_output_weights", False): - output_weight = gpt_model.shared_embedding_or_output_weight() + if getattr(model, "share_embeddings_and_output_weights", False): + output_weight = model.shared_embedding_or_output_weight() if detach_output_weight: # Optionally isolate the shared embedding/output-layer from the draft # gradient too (default keeps NeMo's behaviour: only the trunk hidden @@ -51,6 +65,6 @@ def project_mtp_hidden_to_logits(hidden_states_per_layer, gpt_model, detach_outp logits_per_layer = [] for hidden in hidden_states_per_layer: - logits, _ = gpt_model.output_layer(hidden, weight=output_weight) + logits, _ = model.output_layer(hidden, weight=output_weight) logits_per_layer.append(logits.transpose(0, 1).contiguous()) return logits_per_layer diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index cde0043d14..c9b087b8b9 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -3,11 +3,16 @@ # This is the SkyRL analogue of NeMo-RL's ``HiddenStateCapture`` # (https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/draft/hidden_capture.py). # -# Megatron's ``GPTModel`` runs its native MTP block (``self.mtp``) inside the forward whenever -# ``self.mtp_process`` is set. Crucially, ``MultiTokenPredictionBlock.forward`` passes the trunk -# hidden states through as the *first* chunk of its output, and ``process_mtp_loss`` returns that -# chunk as the hidden states for the *main* logits. So we must NOT tamper with the in-forward call: -# the main policy logits have to stay connected to the trunk for the policy loss to train it. +# Works for any Megatron model that builds native MTP heads and exposes ``self.mtp`` (the shared +# ``MultiTokenPredictionBlock``): ``GPTModel`` (DeepSeek-V3 / GLM / Qwen3-Next, ...) and +# ``MambaModel`` (Qwen3.5 / NemotronH, ...). Both run ``self.mtp(...)`` + ``process_mtp_loss(...)`` +# inside ``forward`` and surface the same output layout, so the capture below is model-agnostic. +# +# Megatron runs its native MTP block (``self.mtp``) inside the forward whenever ``self.mtp_process`` +# is set. Crucially, ``MultiTokenPredictionBlock.forward`` passes the trunk hidden states through as +# a passthrough chunk of its output, and ``process_mtp_loss`` returns that chunk as the hidden states +# for the *main* logits. So we must NOT tamper with the in-forward call: the main policy logits have +# to stay connected to the trunk for the policy loss to train it. # # Instead we: # 1. register a forward pre-hook on ``self.mtp`` that simply *records* the keyword arguments @@ -32,13 +37,29 @@ from typing import List, Optional -def _unwrap_gpt_model(model): - """Return the underlying GPTModel from a (possibly DDP/Float16) wrapper.""" +def _unwrap_model(model): + """Return the underlying Megatron model (GPTModel / MambaModel / ...) from a wrapper.""" from megatron.core.utils import unwrap_model return unwrap_model(model) +def _mtp_layer_offset(mtp_block) -> int: + """Number of passthrough chunks ahead of this stage's MTP-depth chunks. + + ``MultiTokenPredictionBlock.forward`` chunks its input into ``1 + offset`` passthrough chunks + (the trunk hidden states plus any MTP outputs forwarded from earlier pipeline/VP stages), + appends one new chunk per MTP depth, and concatenates everything along dim 0. ``offset`` is 0 in + the common single-stage case; we resolve it via megatron-core so the replay split below is also + correct under PP/VPP. Falls back to 0 if the helper is unavailable (older megatron-core).""" + try: + from megatron.core.transformer.multi_token_prediction import get_mtp_layer_offset + + return int(get_mtp_layer_offset(mtp_block.config, getattr(mtp_block, "vp_stage", None))) + except Exception: + return 0 + + class MTPHiddenCapture: """Record the MTP block's inputs during the forward, then replay it decoupled. @@ -49,7 +70,7 @@ class MTPHiddenCapture: """ def __init__(self, model, detach_trunk: bool = True): - self.gpt_model = _unwrap_gpt_model(model) + self.model = _unwrap_model(model) self.detach_trunk = detach_trunk self._args = None self._kwargs = None @@ -58,7 +79,7 @@ def __init__(self, model, detach_trunk: bool = True): @property def mtp_num_layers(self) -> int: - return int(getattr(self.gpt_model.config, "mtp_num_layers", 0) or 0) + return int(getattr(self.model.config, "mtp_num_layers", 0) or 0) def _pre_hook(self, _module, args, kwargs): # Record (do not modify) the arguments Megatron built for the MTP block. @@ -68,7 +89,7 @@ def _pre_hook(self, _module, args, kwargs): @contextmanager def capture(self): - mtp = getattr(self.gpt_model, "mtp", None) + mtp = getattr(self.model, "mtp", None) if mtp is None: # Model has no MTP heads on this rank/stage; nothing to capture. yield self @@ -106,13 +127,16 @@ def compute_student_hidden_states(self) -> Optional[List]: if hidden is not None and self.detach_trunk: kwargs["hidden_states"] = hidden.detach() - # MultiTokenPredictionBlock returns the concatenated hidden states - # [trunk; mtp_0; mtp_1; ...] along the sequence (dim 0). - captured = self.gpt_model.mtp(*self._args, **kwargs) + mtp = self.model.mtp + # MultiTokenPredictionBlock concatenates, along dim 0: + # [<1 + offset> passthrough chunks (trunk + earlier-stage MTP outputs)] + # + [ new MTP-depth chunks produced on this stage]. + # We want this stage's new MTP-depth chunks (the last `num_layers`). + captured = mtp(*self._args, **kwargs) num_layers = self.mtp_num_layers - chunks = list(torch.chunk(captured, 1 + num_layers, dim=0)) - # Drop the (passthrough) trunk chunk; keep the per-MTP-depth chunks. - return chunks[1:] + total_chunks = 1 + _mtp_layer_offset(mtp) + num_layers + chunks = list(torch.chunk(captured, total_chunks, dim=0)) + return chunks[-num_layers:] @contextmanager diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 682f39cef7..a4ea65bebf 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -1,3 +1,4 @@ +import os from dataclasses import asdict from functools import partial from typing import Any, Callable, Dict, List, Optional @@ -242,11 +243,13 @@ def forward_backward_mini_batch( # ------------------------------------------------------------------ # If the model was built with native MTP heads (policy.megatron_config.mtp_num_layers), # train them with an explicit, decoupled loss instead of Megatron's in-forward - # process_mtp_loss / MTPLossAutoScaler path. We keep the heads running inside GPTModel.forward - # (so we don't have to reconstruct rotary embeddings) but pass NO labels, so process_mtp_loss - # short-circuits and no MTP gradient is coupled onto the trunk. A forward hook captures the - # MTP block's hidden states (with its trunk input optionally detached) and we project + score - # them ourselves. Only active during training (not forward_only / logprob-only passes). + # process_mtp_loss / MTPLossAutoScaler path. This is model-agnostic: any model that exposes a + # native MTP block works (GPTModel for DeepSeek/GLM/Qwen3-Next, MambaModel for Qwen3.5/ + # NemotronH). We keep the heads running inside the model's forward (so we don't have to + # reconstruct rotary embeddings) but pass NO labels, so process_mtp_loss short-circuits and no + # MTP gradient is coupled onto the trunk. A forward hook captures the MTP block's hidden states + # (with its trunk input optionally detached) and we project + score them ourselves. Only + # active during training (not forward_only / logprob-only passes). model_config = get_model_config(self.actor_module[0]) mtp_enabled = (not forward_only) and bool(getattr(model_config, "mtp_num_layers", None)) mcfg = self.cfg.policy.megatron_config @@ -255,6 +258,19 @@ def forward_backward_mini_batch( mtp_detach_trunk = bool(getattr(mcfg, "mtp_detach_trunk", True)) mtp_detach_shared_output = bool(getattr(mcfg, "mtp_detach_shared_output", False)) + if os.environ.get("MTP_DEBUG"): + from megatron.core.utils import unwrap_model as _uw + + _gm = _uw(self.actor_module[0]) + print( + f"[MTP_DEBUG] forward_only={forward_only} mtp_enabled={mtp_enabled} " + f"model_config.mtp_num_layers={getattr(model_config, 'mtp_num_layers', 'MISSING')} " + f"unwrapped_type={type(_gm).__name__} has_mtp_attr={hasattr(_gm, 'mtp')} " + f"mtp_is_none={getattr(_gm, 'mtp', None) is None} " + f"mtp_process={getattr(_gm, 'mtp_process', 'MISSING')}", + flush=True, + ) + # Resolve loss function resolved_loss_name = loss_fn if loss_fn is not None else self.cfg.algorithm.policy_loss_type if loss_fn is not None: @@ -558,9 +574,10 @@ def depad(tensor): ) # Run the policy forward. When MTP is active, a pre-hook records the native MTP block's - # arguments (we pass NO labels, so GPTModel's process_mtp_loss short-circuits and the + # arguments (we pass NO labels, so the model's process_mtp_loss short-circuits and the # main logits stay coupled to the trunk; MTPLossAutoScaler never runs). student_hidden = None + student_model = None with maybe_capture_mtp_hidden(model, mtp_enabled, detach_trunk=mtp_detach_trunk) as capture: outputs = model( new_sequences, @@ -572,7 +589,7 @@ def depad(tensor): # while still inside the capture context (so the MTP block stays in eval mode). if mtp_enabled and capture is not None: student_hidden = capture.compute_student_hidden_states() - gpt_model = capture.gpt_model + student_model = capture.model outputs = depad(outputs) @@ -581,7 +598,7 @@ def depad(tensor): # scored against the policy's own distribution (or hard labels) in loss_func. if student_hidden is not None: student_logits = project_mtp_hidden_to_logits( - student_hidden, gpt_model, detach_output_weight=mtp_detach_shared_output + student_hidden, student_model, detach_output_weight=mtp_detach_shared_output ) batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index ee6ad4ae7d..8eda134d7b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -405,9 +405,11 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) - # Multi-Token Prediction (MTP) head configuration. - # DeepSeek-V3 / GLM-4.7-Flash bridges set provider.mtp_num_layers from the HF config's - # num_nextn_predict_layers. We layer SkyRL's explicit config on top: + # Multi-Token Prediction (MTP) head configuration. Model-agnostic: megatron-bridge populates + # provider.mtp_num_layers from whichever HF key the model uses — DeepSeek/GLM via + # num_nextn_predict_layers, Qwen3.5/Qwen3-Next via the generic + # (mtp_num_hidden_layers -> mtp_num_layers) config mapping. We layer SkyRL's explicit config + # on top: # - enable_mtp=False (ref worker): force MTP off (the heads don't affect ref logprobs # and would only waste compute/memory). # - mtp_num_layers set: override the model default (0 force-disables). diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index f0c2f9c3de..d1aa8dabdb 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -1227,6 +1227,7 @@ def create_ray_wrapped_inference_engines_from_config( "enable_ray_prometheus_stats": cfg.generator.inference_engine.enable_ray_prometheus_stats, "distributed_executor_backend": cfg.generator.inference_engine.distributed_executor_backend, "language_model_only": cfg.generator.inference_engine.language_model_only, + "speculative_config": cfg.generator.inference_engine.speculative_config, } # Conditionally add LoRA parameters if LoRA is enabled diff --git a/skyrl/train/entrypoints/main_base.py b/skyrl/train/entrypoints/main_base.py index 4aa53acac2..d1b41f3d57 100644 --- a/skyrl/train/entrypoints/main_base.py +++ b/skyrl/train/entrypoints/main_base.py @@ -81,6 +81,7 @@ def create_ray_wrapped_inference_engines_from_config( "enable_ray_prometheus_stats": ie_cfg.enable_ray_prometheus_stats, "enable_return_routed_experts": ie_cfg.enable_return_routed_experts, "distributed_executor_backend": ie_cfg.distributed_executor_backend, + "speculative_config": ie_cfg.speculative_config, } # Conditionally add LoRA parameters if LoRA is enabled diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index b9ce0eee73..b60c315303 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -841,6 +841,25 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis return training_input @torch.no_grad() + async def _record_spec_decode_metrics(self) -> None: + """Record the vLLM speculative-decoding (MTP draft) acceptance rate for this rollout step. + + Reads the engines' cumulative draft/accept counters and logs the per-step delta as + ``vllm/draft_acceptance_rate`` (+ raw counts). No-op when speculative decoding is disabled or + the backend doesn't expose the stats. Best-effort: never fails the training step. + """ + from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + acceptance_rate_metrics, + ) + + try: + cumulative = await self.inference_engine_client.get_spec_decode_metrics() + except Exception as e: + logger.warning(f"Failed to read vLLM spec-decode metrics: {e}") + return + metrics, self._prev_spec_decode = acceptance_rate_metrics(cumulative, self._prev_spec_decode) + self.all_metrics.update(metrics) + async def generate( self, input_batch: GeneratorInput, @@ -861,6 +880,9 @@ async def generate( self.all_metrics.update(generator_output["rollout_metrics"]) generator_output.pop("rollout_metrics", None) + # vLLM speculative-decoding (MTP draft) acceptance rate for this rollout step. + await self._record_spec_decode_metrics() + validate_generator_output( len(input_batch["prompts"]), generator_output, diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index c965ccd4c0..446c0fa323 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -218,6 +218,39 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): # TODO (sumanthrh): Most of this should be moved to __post_init__ for the dataclasses +def _apply_mtp_config(cfg: SkyRLTrainConfig): + """Propagate the high-level ``trainer.mtp`` knob to the training + inference configs. + + When MTP is enabled, build/train ``num_speculative_tokens`` native MTP heads (Megatron) with the + decoupled draft loss, and enable vLLM MTP speculative decoding with the same draft depth. When + disabled, force the MTP heads off so they are neither built nor run. + """ + mtp = getattr(cfg.trainer, "mtp", None) + if mtp is None: + return + + mcfg = cfg.trainer.policy.megatron_config + if not mtp.enabled: + # Explicit 0 force-disables MTP even on MTP-capable models. + mcfg.mtp_num_layers = 0 + return + + assert mtp.num_speculative_tokens >= 1, "trainer.mtp.num_speculative_tokens must be >= 1 when enabled" + # Training side: build + train this many native MTP heads with the decoupled draft loss. + mcfg.mtp_num_layers = mtp.num_speculative_tokens + mcfg.mtp_loss_type = mtp.loss_type + mcfg.mtp_loss_weight = mtp.loss_weight + + # Inference side: vLLM MTP speculative decoding with the same draft depth. Don't clobber an + # explicit user-provided speculative_config. + ie_cfg = cfg.generator.inference_engine + if ie_cfg.speculative_config is None: + ie_cfg.speculative_config = { + "method": "mtp", + "num_speculative_tokens": mtp.num_speculative_tokens, + } + + def validate_cfg(cfg: SkyRLTrainConfig): if cfg.trainer.strategy == "fsdp2": import warnings @@ -231,6 +264,12 @@ def validate_cfg(cfg: SkyRLTrainConfig): # Validate generation config separately validate_generator_cfg(cfg) + + # Multi-Token Prediction (MTP): the high-level `trainer.mtp` knob is the single source of truth. + # Propagate it to the training side (Megatron MTP heads + decoupled draft loss) and the inference + # side (vLLM MTP speculative decoding) so both stay consistent. + _apply_mtp_config(cfg) + from skyrl.backends.skyrl_train.utils.ppo_utils import ( AdvantageEstimatorRegistry, PolicyLossRegistry, @@ -676,6 +715,10 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: env_vars["NCCL_P2P_DISABLE"] = "1" env_vars["NCCL_SHM_DISABLE"] = "1" + if os.environ.get("NCCL_NET_PLUGIN"): + logger.info(f"Exporting NCCL_NET_PLUGIN to ray runtime env: {os.environ['NCCL_NET_PLUGIN']}") + env_vars["NCCL_NET_PLUGIN"] = os.environ["NCCL_NET_PLUGIN"] + # TODO: this can be removed if we standardize on env files. # But it's helpful for a quickstart if os.environ.get("WANDB_API_KEY"): @@ -715,6 +758,16 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: logger.info(f"Exporting `SKYRL_RAY_PG_TIMEOUT_IN_S` to ray runtime env: {pg_timeout}") env_vars["SKYRL_RAY_PG_TIMEOUT_IN_S"] = pg_timeout + # Forward uv's project-environment selection to the workers. Ray's uv runtime-env hook makes each + # worker re-run `uv run ... --extra `, and that subprocess must resolve to the SAME venv + # as the driver. Workers are spawned by the raylet and only inherit env vars we forward here, so a + # driver-only `UV_PROJECT_ENVIRONMENT` (e.g. from a local `.env`) would otherwise be lost and the + # worker's `uv run` would fall back to the empty project `.venv` (-> `No module named 'megatron'`). + for var_name in ("UV_PROJECT_ENVIRONMENT", "UV_CACHE_DIR", "UV_LINK_MODE", "UV_PYTHON"): + if value := os.environ.get(var_name): + logger.info(f"Exporting `{var_name}` to ray runtime env: {value}") + env_vars[var_name] = value + return env_vars diff --git a/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py b/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py index 871a151eae..30404cadf3 100644 --- a/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py +++ b/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py @@ -354,6 +354,63 @@ async def generate(self, input_batch): assert out["stop_reasons"][i] == "stop" +# ------------------------------------------- +# tests for InferenceEngineClient.get_spec_decode_metrics +# -------------------------------------------- + + +def _make_spec_client(per_engine_stats): + """Build a client whose engines report the given per-engine spec-decode stats.""" + + class MockEngine: + def __init__(self, stats): + self._stats = stats + + def dp_size(self): + return 1 + + async def get_spec_decode_metrics(self): + return self._stats + + cfg = _make_min_cfg() + return InferenceEngineClient( + engines=[MockEngine(s) for s in per_engine_stats], + tokenizer=object(), + model_path=cfg.trainer.policy.model.path, + lora_cfg=cfg.trainer.policy.model.lora, + inference_engine_cfg=cfg.generator.inference_engine, + ) + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_sums_across_engines(): + client = _make_spec_client( + [ + {"num_drafts": 3, "num_draft_tokens": 10, "num_accepted_tokens": 6}, + {"num_drafts": 1, "num_draft_tokens": 4, "num_accepted_tokens": 2}, + ] + ) + totals = await client.get_spec_decode_metrics() + assert totals == {"num_drafts": 4, "num_draft_tokens": 14, "num_accepted_tokens": 8} + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_returns_none_when_no_engine_reports(): + # Speculative decoding disabled / unsupported on every engine -> None (not a zero-dict). + client = _make_spec_client([None, None]) + assert await client.get_spec_decode_metrics() is None + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_ignores_non_reporting_engines(): + # A mix of reporting and non-reporting engines still aggregates the reporters only. + client = _make_spec_client( + [None, {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5}] + ) + totals = await client.get_spec_decode_metrics() + assert totals == {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5} + + # ----------------------------- # Test for route_prompts_to_engines function that routes prompts to inference engines # in inference engine client. diff --git a/tests/backends/skyrl_train/mtp/test_hidden_capture.py b/tests/backends/skyrl_train/mtp/test_hidden_capture.py index 634ddff633..edef1aae0f 100644 --- a/tests/backends/skyrl_train/mtp/test_hidden_capture.py +++ b/tests/backends/skyrl_train/mtp/test_hidden_capture.py @@ -2,7 +2,9 @@ These use a fake MTP block (no Megatron) to verify the capture records the block's arguments and that the decoupled replay detaches the trunk hidden states so the -draft gradient never reaches the policy backbone. +draft gradient never reaches the policy backbone. The capture is model-agnostic +(GPTModel for DeepSeek/GLM/Qwen3-Next, MambaModel for Qwen3.5/NemotronH); the fakes +below stand in for either base class. uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_hidden_capture.py """ @@ -14,44 +16,70 @@ import torch.nn as nn # Stub out megatron.core.utils.unwrap_model so hidden_capture imports on CPU. +# (hidden_capture also tries megatron.core.transformer.multi_token_prediction.get_mtp_layer_offset, +# which is absent here, so it falls back to offset 0 — the single-stage / no-PP case.) _fake_mcore_utils = types.ModuleType("megatron.core.utils") _fake_mcore_utils.unwrap_model = lambda m: m sys.modules.setdefault("megatron", types.ModuleType("megatron")) sys.modules.setdefault("megatron.core", types.ModuleType("megatron.core")) sys.modules["megatron.core.utils"] = _fake_mcore_utils +from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits # noqa: E402 from skyrl.backends.skyrl_train.mtp.hidden_capture import MTPHiddenCapture # noqa: E402 class _FakeMTPBlock(nn.Module): - """Mimics MultiTokenPredictionBlock: returns cat([trunk; mtp_0], dim=0).""" + """Mimics MultiTokenPredictionBlock: returns cat([trunk; mtp_0; ...; mtp_{k-1}], dim=0).""" - def __init__(self, hidden): + def __init__(self, hidden, num_layers=1): super().__init__() - self.w = nn.Parameter(torch.ones(hidden)) + self.num_layers = num_layers + # One distinct weight per MTP depth so each depth's output is separable. + self.w = nn.ParameterList([nn.Parameter(torch.ones(hidden) * (i + 1)) for i in range(num_layers)]) def forward(self, hidden_states, **kwargs): - mtp_0 = hidden_states * self.w # depends on params AND trunk hidden - return torch.cat([hidden_states, mtp_0], dim=0) + # depends on params AND trunk hidden, like the real block. + chunks = [hidden_states] + [hidden_states * self.w[i] for i in range(self.num_layers)] + return torch.cat(chunks, dim=0) class _FakeGPT(nn.Module): + """Stands in for any Megatron model exposing a native ``.mtp`` block + ``.config``.""" + def __init__(self, hidden=4, mtp_num_layers=1): super().__init__() - self.mtp = _FakeMTPBlock(hidden) + self.mtp = _FakeMTPBlock(hidden, num_layers=mtp_num_layers) self.config = types.SimpleNamespace(mtp_num_layers=mtp_num_layers) -def _run(detach_trunk): - gpt = _FakeGPT() - capture = MTPHiddenCapture(gpt, detach_trunk=detach_trunk) +class _FakeMamba(_FakeGPT): + """Stands in for MambaModel: same MTP surface plus the shared-output-layer surface the + adapter's projection uses (``output_layer`` + tied ``shared_embedding_or_output_weight``).""" + + def __init__(self, hidden=4, vocab=5, mtp_num_layers=1): + super().__init__(hidden=hidden, mtp_num_layers=mtp_num_layers) + self.share_embeddings_and_output_weights = True + self._out_weight = nn.Parameter(torch.randn(vocab, hidden)) + + def shared_embedding_or_output_weight(self): + return self._out_weight + + def output_layer(self, hidden, weight=None): + w = weight if weight is not None else self._out_weight + # [seq, batch, hidden] @ [hidden, vocab] -> [seq, batch, vocab]; mirror Megatron's (logits, bias). + return hidden @ w.t(), None + + +def _run(detach_trunk, model_cls=_FakeGPT, mtp_num_layers=1): + model = model_cls(mtp_num_layers=mtp_num_layers) + capture = MTPHiddenCapture(model, detach_trunk=detach_trunk) s, b, h = 3, 2, 4 trunk = torch.randn(s, b, h, requires_grad=True) with capture.capture(): - # Simulate GPTModel's in-forward MTP call (records kwargs via the pre-hook). - _ = gpt.mtp(hidden_states=trunk, position_ids=torch.zeros(b, s)) + # Simulate the model's in-forward MTP call (records kwargs via the pre-hook). + _ = model.mtp(hidden_states=trunk, position_ids=torch.zeros(b, s)) student = capture.compute_student_hidden_states() - return gpt, trunk, student + return model, trunk, student def test_capture_returns_one_chunk_per_mtp_depth(): @@ -60,27 +88,66 @@ def test_capture_returns_one_chunk_per_mtp_depth(): assert student[0].shape == (3, 2, 4) +def test_capture_returns_one_chunk_per_depth_multilayer(): + # chunk[-num_layers:] must return exactly this stage's MTP-depth outputs. + _, _, student = _run(detach_trunk=True, mtp_num_layers=2) + assert student is not None and len(student) == 2 + assert all(chunk.shape == (3, 2, 4) for chunk in student) + + def test_replay_detaches_trunk_when_detach_trunk_true(): - gpt, trunk, student = _run(detach_trunk=True) + model, trunk, student = _run(detach_trunk=True) student[0].sum().backward() # The MTP-head parameter receives gradient... - assert gpt.mtp.w.grad is not None and gpt.mtp.w.grad.abs().sum() > 0 + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 # ...but the trunk hidden states do NOT (decoupled). assert trunk.grad is None def test_replay_couples_trunk_when_detach_trunk_false(): - gpt, trunk, student = _run(detach_trunk=False) + model, trunk, student = _run(detach_trunk=False) student[0].sum().backward() - assert gpt.mtp.w.grad is not None and gpt.mtp.w.grad.abs().sum() > 0 + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 # With detach disabled, the gradient flows back into the trunk hidden states. assert trunk.grad is not None and trunk.grad.abs().sum() > 0 +def test_capture_is_model_agnostic_mambamodel(): + # The capture must work identically for a MambaModel-like base (Qwen3.5), exposing the same + # .mtp surface. capture.model is the unwrapped model (renamed from the old GPTModel-specific attr). + model, trunk, student = _run(detach_trunk=True, model_cls=_FakeMamba) + assert isinstance(MTPHiddenCapture(model).model, _FakeMamba) + assert student is not None and len(student) == 1 + assert student[0].shape == (3, 2, 4) + + +def test_project_to_logits_decoupled_end_to_end(): + # End-to-end through the adapter: capture (detached) -> shared output layer -> student logits. + # The draft gradient reaches the MTP head + shared output weight, never the trunk. + model, trunk, student = _run(detach_trunk=True, model_cls=_FakeMamba) + logits = project_mtp_hidden_to_logits(student, model) + assert len(logits) == 1 + # [seq, batch, vocab] -> transposed to [batch, seq, vocab]. + assert logits[0].shape == (2, 3, 5) + logits[0].sum().backward() + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 + assert model._out_weight.grad is not None and model._out_weight.grad.abs().sum() > 0 + assert trunk.grad is None + + +def test_project_to_logits_detach_shared_output(): + # detach_output_weight=True isolates the shared embedding/output weight from the draft gradient. + model, trunk, student = _run(detach_trunk=True, model_cls=_FakeMamba) + logits = project_mtp_hidden_to_logits(student, model, detach_output_weight=True) + logits[0].sum().backward() + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 + assert model._out_weight.grad is None + + def test_no_capture_when_block_absent(): - gpt = _FakeGPT() - gpt.mtp = None - capture = MTPHiddenCapture(gpt, detach_trunk=True) + model = _FakeGPT() + model.mtp = None + capture = MTPHiddenCapture(model, detach_trunk=True) with capture.capture(): pass assert capture.compute_student_hidden_states() is None diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py index a69c3bf07c..4d2c15408d 100644 --- a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py @@ -3,14 +3,13 @@ uv run --isolated --extra dev --extra megatron pytest tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py """ -import asyncio import types from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + acceptance_rate_metrics, make_spec_decode_stat_logger_class, sum_spec_decode_loggers, ) -from skyrl.train.generators.skyrl_gym_generator import SkyRLGymGenerator def _sched(num_draft, num_accepted, num_drafts): @@ -44,46 +43,34 @@ def test_sum_spec_decode_loggers(): assert sum_spec_decode_loggers([]) is None -class _FakeClient: - def __init__(self, values): - self._values = list(values) - self._i = 0 - - async def get_spec_decode_metrics(self): - v = self._values[min(self._i, len(self._values) - 1)] - self._i += 1 - return v - - -def _make_generator(client): - gen = SkyRLGymGenerator.__new__(SkyRLGymGenerator) - gen.inference_engine_client = client - gen._prev_spec_decode = None - return gen - - -def test_generator_acceptance_rate_per_step_delta(): +def test_acceptance_rate_per_step_delta(): # cumulative counters grow each step; the metric is the per-step delta ratio. - client = _FakeClient( - [ - {"num_drafts": 5, "num_draft_tokens": 10, "num_accepted_tokens": 6}, - {"num_drafts": 11, "num_draft_tokens": 22, "num_accepted_tokens": 15}, - ] + prev = None + m1, prev = acceptance_rate_metrics( + {"num_drafts": 5, "num_draft_tokens": 10, "num_accepted_tokens": 6}, prev ) - gen = _make_generator(client) - - m1 = asyncio.run(gen._spec_decode_rollout_metrics()) assert m1["vllm/draft_num_draft_tokens"] == 10 assert m1["vllm/draft_num_accepted_tokens"] == 6 assert abs(m1["vllm/draft_acceptance_rate"] - 0.6) < 1e-9 - m2 = asyncio.run(gen._spec_decode_rollout_metrics()) + m2, prev = acceptance_rate_metrics( + {"num_drafts": 11, "num_draft_tokens": 22, "num_accepted_tokens": 15}, prev + ) # deltas: drafted 22-10=12, accepted 15-6=9 -> 0.75 assert m2["vllm/draft_num_draft_tokens"] == 12 assert m2["vllm/draft_num_accepted_tokens"] == 9 assert abs(m2["vllm/draft_acceptance_rate"] - 0.75) < 1e-9 -def test_generator_no_spec_decode_returns_empty(): - gen = _make_generator(_FakeClient([None])) - assert asyncio.run(gen._spec_decode_rollout_metrics()) == {} +def test_acceptance_rate_no_spec_decode_returns_empty(): + metrics, prev = acceptance_rate_metrics(None, None) + assert metrics == {} + assert prev is None + + +def test_acceptance_rate_no_drafts_omits_rate(): + # zero drafted tokens -> report counts but no (undefined) rate. + metrics, _ = acceptance_rate_metrics( + {"num_drafts": 0, "num_draft_tokens": 0, "num_accepted_tokens": 0}, None + ) + assert metrics == {"vllm/draft_num_draft_tokens": 0, "vllm/draft_num_accepted_tokens": 0} diff --git a/tests/train/test_mtp_config.py b/tests/train/test_mtp_config.py index 1dcf46fb0a..49ad606e7d 100644 --- a/tests/train/test_mtp_config.py +++ b/tests/train/test_mtp_config.py @@ -88,3 +88,24 @@ def test_apply_mtp_config_does_not_clobber_explicit_speculative_config(): cfg.generator.inference_engine.speculative_config = {"method": "mtp", "num_speculative_tokens": 5} _apply_mtp_config(cfg) assert cfg.generator.inference_engine.speculative_config["num_speculative_tokens"] == 5 + + +def test_local_engine_builder_accepts_speculative_config(): + # Regression: the colocated local-engine path must accept + forward speculative_config to vLLM. + # It was only wired into the standalone HTTP-server path (build_vllm_cli_args), so colocated runs + # (the common MTP path) silently ran with spec decode OFF -> 0 drafted / 0 accepted tokens. + # + # Kept lightweight on purpose: we assert the engine builder *accepts* a speculative_config kwarg + # (the parameter that was missing), rather than calling the from-config wrapper. The wrapper lives + # in main_base, whose import chain (trainer -> ppo_utils) does a Ray auto-init at import time and + # would touch the live cluster under the session-scoped ray fixture. + import inspect + + from skyrl.backends.skyrl_train.inference_engines.ray_wrapped_inference_engine import ( + create_ray_wrapped_inference_engines, + ) + + params = inspect.signature(create_ray_wrapped_inference_engines).parameters + assert "speculative_config" in params, "engine builder must accept speculative_config (MTP spec decode)" + # Default must be None so non-spec-decode runs are unaffected. + assert params["speculative_config"].default is None From 99b65aa3ac6faec1bcfde1fd87ad6cca7475b805 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 8 Jun 2026 03:22:57 +0000 Subject: [PATCH 06/28] update model capture for Mamba Models --- .../skyrl_train/mtp/hidden_capture.py | 27 ++++++++++++++++++- .../megatron/megatron_model_wrapper.py | 12 ++++++--- skyrl/train/utils/utils.py | 9 ++++++- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index c9b087b8b9..43f85fb0cb 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -44,6 +44,27 @@ def _unwrap_model(model): return unwrap_model(model) +def _resolve_mtp_host(model): + """Return the submodule that actually owns the native MTP block (``.mtp``). + + For a plain ``GPTModel`` / ``MambaModel`` that is the model itself. For vision-language wrappers + (e.g. Qwen3.5-VL's ``Qwen3VLModel``) the text backbone *and* its MTP head are nested one level + down at ``model.language_model`` (the bridge maps the heads to ``language_model.mtp.*``), so the + top-level model has no ``.mtp`` and the capture would otherwise silently find nothing. We descend + through the common ``.language_model`` nesting until we find a module exposing ``.mtp``; if none + does, return the top-level model unchanged (capture then yields ``None``, as before).""" + seen = set() + cur = model + for _ in range(4): # bounded descent guard against pathological nesting / cycles + if cur is None or id(cur) in seen: + break + seen.add(id(cur)) + if getattr(cur, "mtp", None) is not None: + return cur + cur = getattr(cur, "language_model", None) + return model + + def _mtp_layer_offset(mtp_block) -> int: """Number of passthrough chunks ahead of this stage's MTP-depth chunks. @@ -70,7 +91,11 @@ class MTPHiddenCapture: """ def __init__(self, model, detach_trunk: bool = True): - self.model = _unwrap_model(model) + # Resolve the module that owns the MTP block: the unwrapped model itself for plain + # GPT/Mamba models, or the nested ``.language_model`` for VL wrappers (e.g. Qwen3.5-VL). + # capture.model is also what the caller projects through (output_layer / shared weight), + # so it must be the same module that holds the heads. + self.model = _resolve_mtp_host(_unwrap_model(model)) self.detach_trunk = detach_trunk self._args = None self._kwargs = None diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index a4ea65bebf..a64da1ea69 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -261,13 +261,19 @@ def forward_backward_mini_batch( if os.environ.get("MTP_DEBUG"): from megatron.core.utils import unwrap_model as _uw + from skyrl.backends.skyrl_train.mtp.hidden_capture import _resolve_mtp_host + _gm = _uw(self.actor_module[0]) + _host = _resolve_mtp_host(_gm) + _lm = getattr(_gm, "language_model", None) print( f"[MTP_DEBUG] forward_only={forward_only} mtp_enabled={mtp_enabled} " f"model_config.mtp_num_layers={getattr(model_config, 'mtp_num_layers', 'MISSING')} " - f"unwrapped_type={type(_gm).__name__} has_mtp_attr={hasattr(_gm, 'mtp')} " - f"mtp_is_none={getattr(_gm, 'mtp', None) is None} " - f"mtp_process={getattr(_gm, 'mtp_process', 'MISSING')}", + f"unwrapped_type={type(_gm).__name__} top_has_mtp={getattr(_gm, 'mtp', None) is not None} " + f"has_language_model={_lm is not None} " + f"lm_has_mtp={getattr(_lm, 'mtp', None) is not None if _lm is not None else 'NA'} " + f"resolved_host={type(_host).__name__} host_mtp_is_none={getattr(_host, 'mtp', None) is None} " + f"host_mtp_process={getattr(_host, 'mtp_process', 'MISSING')}", flush=True, ) diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 446c0fa323..24efcb8987 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -763,7 +763,14 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: # as the driver. Workers are spawned by the raylet and only inherit env vars we forward here, so a # driver-only `UV_PROJECT_ENVIRONMENT` (e.g. from a local `.env`) would otherwise be lost and the # worker's `uv run` would fall back to the empty project `.venv` (-> `No module named 'megatron'`). - for var_name in ("UV_PROJECT_ENVIRONMENT", "UV_CACHE_DIR", "UV_LINK_MODE", "UV_PYTHON"): + for var_name in ( + "UV_PROJECT_ENVIRONMENT", + "UV_CACHE_DIR", + "UV_LINK_MODE", + "UV_PYTHON", + "MTP_DEBUG", + "PYTORCH_CUDA_ALLOC_CONF", + ): if value := os.environ.get(var_name): logger.info(f"Exporting `{var_name}` to ray runtime env: {value}") env_vars[var_name] = value From 64799a0d5f231d21dcc45f780e0ae8929bf2f93a Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 8 Jun 2026 03:27:00 +0000 Subject: [PATCH 07/28] reorganize tests, improve CE loss to be more memory efficient --- skyrl/backends/skyrl_train/mtp/soft_ce.py | 87 +++++++++++-------- .../test_build_vllm_cli_args.py | 21 +---- .../mtp/test_build_vllm_cli_args_mtp.py | 31 +++++++ .../skyrl_train/mtp/test_hidden_capture.py | 46 ++++++++++ .../skyrl_train/mtp}/test_mtp_config.py | 0 .../{utils => mtp}/test_mtp_torch_utils.py | 0 .../backends/skyrl_train/mtp/test_soft_ce.py | 41 +++++++++ 7 files changed, 169 insertions(+), 57 deletions(-) create mode 100644 tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py rename tests/{train => backends/skyrl_train/mtp}/test_mtp_config.py (100%) rename tests/backends/skyrl_train/{utils => mtp}/test_mtp_torch_utils.py (100%) diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index 05896d55c2..d0025ecf5b 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -61,52 +61,63 @@ def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.T return rolled +def _vocab_parallel_softmax(vocab_parallel_logits, group): + """Global softmax over a TP-sharded vocab dim, using in-place ops to avoid extra full-vocab + allocations. Mirrors NeMo-RL's ``_compute_distributed_softmax``.""" + logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) + torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) + # `- logits_max` allocates once (frees the input copy); everything after is in place. + exp_logits = (vocab_parallel_logits - logits_max).exp_() + sum_exp = exp_logits.sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=group) + return exp_logits.div_(sum_exp) + + +def _vocab_parallel_log_softmax(vocab_parallel_logits, group): + """Global log-softmax over a TP-sharded vocab dim. Mirrors NeMo-RL's + ``_compute_distributed_log_softmax``; the single ``.exp()`` temporary is summed and freed.""" + logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) + torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) + shifted = vocab_parallel_logits - logits_max + sum_exp = shifted.exp().sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=group) + return shifted.sub_(sum_exp.log_()) + + class _VocabParallelSoftCrossEntropy(torch.autograd.Function): - """Soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` - when ``student``/``teacher`` logits are sharded across the vocab (TP) dim. - - Forward reduces across the tensor-parallel group so the softmax denominators - and the teacher/student dot product are computed over the *global* vocab. - Backward returns the classic cross-entropy gradient - ``softmax(student) - softmax(teacher)`` (the teacher is treated as a - constant; no gradient flows to it). This is the TP analogue of NeMo-RL's - ``DistributedCrossEntropy``. + """Soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` when + ``student``/``teacher`` logits are sharded across the vocab (TP) dim. + + Memory-lean port of NeMo-RL's ``DistributedCrossEntropy`` (``nemo_rl/distributed/model_utils.py``): + the per-token cross-entropy is a single ``einsum`` (no full-vocab product tensor) and the student + log-prob buffer is reused **in place** for the backward softmax, so only two full-vocab tensors + are kept (vs ~6-8 in the naive form -- the difference is several GiB at a 248K vocab). Forward + all-reduces across the TP group so the softmax normalizers and the teacher/student dot are over + the *global* vocab. Backward returns the classic cross-entropy gradient + ``softmax(student) - softmax(teacher)`` (teacher is a detached target; no gradient flows to it). """ @staticmethod def forward(ctx, student_vp_logits, teacher_vp_logits, tp_group): - student = student_vp_logits.float() - teacher = teacher_vp_logits.float() - - def global_softmax(logits): - local_max = logits.max(dim=-1, keepdim=True).values - torch.distributed.all_reduce(local_max, op=torch.distributed.ReduceOp.MAX, group=tp_group) - shifted = logits - local_max - exp = shifted.exp() - sum_exp = exp.sum(dim=-1, keepdim=True) - torch.distributed.all_reduce(sum_exp, group=tp_group) - log_sum_exp = local_max + sum_exp.log() - return exp / sum_exp, log_sum_exp - - student_softmax, student_lse = global_softmax(student) - teacher_softmax, _ = global_softmax(teacher) - - # dot = sum_v teacher_softmax_v * student_logit_v (global over vocab) - dot = (teacher_softmax * student).sum(dim=-1, keepdim=True) - torch.distributed.all_reduce(dot, group=tp_group) - - # soft CE = sum_v teacher_p_v * (log_sum_exp - student_logit_v) - # = log_sum_exp - dot (since sum_v teacher_p_v = 1) - per_token_loss = (student_lse - dot).squeeze(-1) - - ctx.save_for_backward(student_softmax, teacher_softmax) - return per_token_loss + ctx.input_dtype = student_vp_logits.dtype + # Detached teacher target distribution (global softmax over the sharded vocab). + target_probs = _vocab_parallel_softmax(teacher_vp_logits.float(), tp_group) + # Student global log-probs; the same buffer is turned into softmax(student) in place below. + student_log_probs = _vocab_parallel_log_softmax(student_vp_logits.float(), tp_group) + # soft CE = -sum_v p_teacher_v * log q_student_v: dot over the local shard, then reduce. + per_token_loss = torch.einsum("...v,...v->...", target_probs, student_log_probs).neg_() + torch.distributed.all_reduce(per_token_loss, op=torch.distributed.ReduceOp.SUM, group=tp_group) + # exp_ in place: log_softmax(student) -> softmax(student) for the gradient (no new alloc). + student_probs = student_log_probs.exp_() + ctx.save_for_backward(student_probs, target_probs) + return per_token_loss.contiguous() @staticmethod def backward(ctx, grad_output): - student_softmax, teacher_softmax = ctx.saved_tensors - grad_student = (student_softmax - teacher_softmax) * grad_output.unsqueeze(-1) - return grad_student.to(student_softmax.dtype), None, None + student_probs, target_probs = ctx.saved_tensors + # d(H(p, q))/d(student_logit_v) = softmax(student)_v - softmax(teacher)_v + grad_student = (student_probs - target_probs) * grad_output.unsqueeze(-1) + return grad_student.to(ctx.input_dtype), None, None def draft_soft_ce( diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py index 650a3464c2..ea0fa85eb2 100644 --- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py +++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py @@ -27,22 +27,5 @@ def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): assert args.tensor_parallel_size == cfg.generator.inference_engine.tensor_parallel_size assert vllm.platforms.current_platform.device_type == "cuda" - -@pytest.mark.vllm -def test_build_vllm_cli_args_speculative_config_mtp(monkeypatch): - """speculative_config is passed through to vLLM for MTP draft decoding.""" - import vllm.platforms - from vllm.platforms.interface import UnspecifiedPlatform - - monkeypatch.setattr(vllm.platforms, "_current_platform", UnspecifiedPlatform()) - - cfg = SkyRLTrainConfig() - # Default: no speculative decoding. - args = build_vllm_cli_args(cfg) - assert getattr(args, "speculative_config", None) is None - - # Enable MTP speculative decoding. - spec = {"method": "mtp", "num_speculative_tokens": 1} - cfg.generator.inference_engine.speculative_config = spec - args = build_vllm_cli_args(cfg) - assert args.speculative_config == spec + # NOTE: the MTP speculative_config wiring test lives in + # tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py diff --git a/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py b/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py new file mode 100644 index 0000000000..d3f0bdc760 --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py @@ -0,0 +1,31 @@ +"""MTP speculative-decoding wiring into the vLLM CLI args (runs on a GPU-less host). + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py +""" + +import pytest + +from skyrl.backends.skyrl_train.inference_servers.utils import build_vllm_cli_args +from skyrl.train.config import SkyRLTrainConfig + + +@pytest.mark.vllm +def test_build_vllm_cli_args_speculative_config_mtp(monkeypatch): + """speculative_config is passed through to vLLM for MTP draft decoding.""" + import vllm.platforms + from vllm.platforms.interface import UnspecifiedPlatform + + # Pin the platform so add_cli_args works on a GPU-less host (see the sibling + # test_build_vllm_cli_args.py for the rationale). + monkeypatch.setattr(vllm.platforms, "_current_platform", UnspecifiedPlatform()) + + cfg = SkyRLTrainConfig() + # Default: no speculative decoding. + args = build_vllm_cli_args(cfg) + assert getattr(args, "speculative_config", None) is None + + # Enable MTP speculative decoding. + spec = {"method": "mtp", "num_speculative_tokens": 1} + cfg.generator.inference_engine.speculative_config = spec + args = build_vllm_cli_args(cfg) + assert args.speculative_config == spec diff --git a/tests/backends/skyrl_train/mtp/test_hidden_capture.py b/tests/backends/skyrl_train/mtp/test_hidden_capture.py index edef1aae0f..59a98fa569 100644 --- a/tests/backends/skyrl_train/mtp/test_hidden_capture.py +++ b/tests/backends/skyrl_train/mtp/test_hidden_capture.py @@ -144,6 +144,52 @@ def test_project_to_logits_detach_shared_output(): assert model._out_weight.grad is None +class _FakeVL(nn.Module): + """Stands in for a vision-language wrapper (e.g. Qwen3.5-VL ``Qwen3VLModel``): the text backbone + and its MTP head are nested at ``.language_model``, and the top-level model has no ``.mtp``.""" + + def __init__(self, hidden=4, vocab=5, mtp_num_layers=1): + super().__init__() + self.language_model = _FakeMamba(hidden=hidden, vocab=vocab, mtp_num_layers=mtp_num_layers) + # Deliberately NO top-level .mtp / .config — mirrors the real VL wrapper. + + +def test_capture_descends_into_language_model_for_vl_wrapper(): + # Regression: Qwen3.5-VL nests the GPTModel + MTP head at model.language_model.mtp. The capture + # must resolve that host instead of looking for a (nonexistent) top-level .mtp, otherwise no + # student hidden states are produced and the draft loss is silently never computed. + from skyrl.backends.skyrl_train.mtp.hidden_capture import _resolve_mtp_host + + model = _FakeVL(mtp_num_layers=1) + assert getattr(model, "mtp", None) is None # top-level has no MTP block... + host = _resolve_mtp_host(model) + assert host is model.language_model # ...capture resolves the nested language model. + + capture = MTPHiddenCapture(model, detach_trunk=True) + assert capture.model is model.language_model + assert capture.mtp_num_layers == 1 + + s, b, h = 3, 2, 4 + trunk = torch.randn(s, b, h, requires_grad=True) + with capture.capture(): + _ = model.language_model.mtp(hidden_states=trunk, position_ids=torch.zeros(b, s)) + student = capture.compute_student_hidden_states() + assert student is not None and len(student) == 1 + assert student[0].shape == (3, 2, 4) + # End-to-end projection through the nested language model's shared output layer. + logits = project_mtp_hidden_to_logits(student, capture.model) + assert logits[0].shape == (2, 3, 5) + + +def test_resolve_mtp_host_returns_model_when_no_mtp_anywhere(): + from skyrl.backends.skyrl_train.mtp.hidden_capture import _resolve_mtp_host + + bare = _FakeGPT() + bare.mtp = None + # No .mtp on the model and no .language_model nesting -> return the model unchanged. + assert _resolve_mtp_host(bare) is bare + + def test_no_capture_when_block_absent(): model = _FakeGPT() model.mtp = None diff --git a/tests/train/test_mtp_config.py b/tests/backends/skyrl_train/mtp/test_mtp_config.py similarity index 100% rename from tests/train/test_mtp_config.py rename to tests/backends/skyrl_train/mtp/test_mtp_config.py diff --git a/tests/backends/skyrl_train/utils/test_mtp_torch_utils.py b/tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py similarity index 100% rename from tests/backends/skyrl_train/utils/test_mtp_torch_utils.py rename to tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py diff --git a/tests/backends/skyrl_train/mtp/test_soft_ce.py b/tests/backends/skyrl_train/mtp/test_soft_ce.py index 61babf2698..adac63fdb7 100644 --- a/tests/backends/skyrl_train/mtp/test_soft_ce.py +++ b/tests/backends/skyrl_train/mtp/test_soft_ce.py @@ -18,6 +18,47 @@ ) +def test_vocab_parallel_soft_ce_matches_reference(monkeypatch): + # The memory-lean _VocabParallelSoftCrossEntropy (NeMo-RL-style einsum + in-place softmax) must + # match the plain full-vocab soft CE in both forward and gradient. We stub the TP all-reduce to a + # no-op so a single shard behaves like the full (un-sharded) vocab, exercising the kernel on CPU. + import torch.distributed as dist + + from skyrl.backends.skyrl_train.mtp.soft_ce import _VocabParallelSoftCrossEntropy + + monkeypatch.setattr(dist, "all_reduce", lambda t, op=None, group=None: t) + + torch.manual_seed(0) + student = torch.randn(2, 4, 7, requires_grad=True) + teacher = torch.randn(2, 4, 7) + g_out = torch.randn(2, 4) + + loss = _VocabParallelSoftCrossEntropy.apply(student, teacher, object()) + loss.backward(g_out) + got_loss, got_grad = loss.detach(), student.grad.detach() + + student2 = student.detach().clone().requires_grad_(True) + ref = -(F.softmax(teacher, -1) * F.log_softmax(student2, -1)).sum(-1) + ref.backward(g_out) + + assert torch.allclose(got_loss, ref.detach(), atol=1e-5) + assert torch.allclose(got_grad, student2.grad, atol=1e-5) + + +def test_vocab_parallel_soft_ce_preserves_input_dtype(monkeypatch): + # Backward must return a grad in the student logits' original dtype (e.g. bf16), not fp32. + import torch.distributed as dist + + from skyrl.backends.skyrl_train.mtp.soft_ce import _VocabParallelSoftCrossEntropy + + monkeypatch.setattr(dist, "all_reduce", lambda t, op=None, group=None: t) + + student = torch.randn(1, 3, 5, dtype=torch.bfloat16, requires_grad=True) + teacher = torch.randn(1, 3, 5, dtype=torch.bfloat16) + _VocabParallelSoftCrossEntropy.apply(student, teacher, object()).sum().backward() + assert student.grad.dtype == torch.bfloat16 + + def test_soft_ce_matches_reference(): torch.manual_seed(0) student = torch.randn(2, 5, 7, requires_grad=True) From 98202fe01f7b48c8a5fe46203e477dd270363156 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 8 Jun 2026 20:45:30 +0000 Subject: [PATCH 08/28] fix bug, mtp storing student logits across microbatches, causing OOM --- .../megatron/megatron_model_wrapper.py | 44 +++++++++++++++---- 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index a64da1ea69..9926364202 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -36,6 +36,7 @@ build_teacher_logits, draft_hard_ce, draft_soft_ce, + draft_soft_ce_topk, shift_mask_for_mtp, ) from skyrl.backends.skyrl_train.utils.torch_utils import ( @@ -257,6 +258,8 @@ def forward_backward_mini_batch( mtp_loss_weight = float(getattr(mcfg, "mtp_loss_weight", 0.1)) mtp_detach_trunk = bool(getattr(mcfg, "mtp_detach_trunk", True)) mtp_detach_shared_output = bool(getattr(mcfg, "mtp_detach_shared_output", False)) + mtp_loss_chunk_size = getattr(mcfg, "mtp_loss_chunk_size", 1024) + mtp_loss_topk = getattr(mcfg, "mtp_loss_topk", None) if os.environ.get("MTP_DEBUG"): from megatron.core.utils import unwrap_model as _uw @@ -361,19 +364,42 @@ def loss_func(logits, data): layer_mask, vocab_parallel_group=tp_grp, vocab_start_index=tp_rank * vocab_size_tp, + chunk_size=mtp_loss_chunk_size, ) ) else: - teacher_logits = build_teacher_logits(teacher_src, layer_idx, detach=True) - per_layer_losses.append( - draft_soft_ce( - student_logits, - teacher_logits, - layer_mask, - vocab_parallel_group=tp_grp, + if mtp_loss_topk: + # Top-k draft loss: O(seq*k) memory, no full-vocab softmax (fits + avoids + # fragmentation at large vocab). Reconciled across the TP group, so it + # scales to any tensor-parallel size incl. cross-node. Pass the *un-rolled* + # policy logits + roll_shift so top-k runs on the policy's own logits (no + # ~[S, vocab] rolled-teacher copy); only the small [B, S, k] top-k is rolled. + per_layer_losses.append( + draft_soft_ce_topk( + student_logits, + teacher_src, + layer_mask, + k=mtp_loss_topk, + vocab_parallel_group=tp_grp, + roll_shift=layer_idx + 1, + ) ) - ) - draft_loss = torch.stack(per_layer_losses).mean() + else: + teacher_logits = build_teacher_logits(teacher_src, layer_idx, detach=True) + per_layer_losses.append( + draft_soft_ce( + student_logits, + teacher_logits, + layer_mask, + vocab_parallel_group=tp_grp, + chunk_size=mtp_loss_chunk_size, + ) + ) + draft_loss = torch.stack(per_layer_losses).mean() + # Drop the dict's reference, this microbatch's autograd graph still holds the + # tensor for its own backward, after which it is freed instead of lingering. + del data["mtp_student_logits"] + student_logits_list = None # SFT path: cross_entropy loss (negative log likelihood) if resolved_loss_name == "cross_entropy": From 7ba31fc4842276bf060292afa9dd653c922dbb2a Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 8 Jun 2026 20:46:09 +0000 Subject: [PATCH 09/28] add top k soft CE loss --- skyrl/backends/skyrl_train/mtp/soft_ce.py | 241 +++++++++++++++--- skyrl/train/config/config.py | 15 ++ skyrl/train/utils/utils.py | 1 + .../test_inference_engine_client.py | 57 +---- .../backends/skyrl_train/mtp/test_soft_ce.py | 120 +++++++++ .../mtp/test_spec_decode_client.py | 89 +++++++ 6 files changed, 433 insertions(+), 90 deletions(-) create mode 100644 tests/backends/skyrl_train/mtp/test_spec_decode_client.py diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index d0025ecf5b..ffb70c04b2 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -62,26 +62,31 @@ def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.T def _vocab_parallel_softmax(vocab_parallel_logits, group): - """Global softmax over a TP-sharded vocab dim, using in-place ops to avoid extra full-vocab - allocations. Mirrors NeMo-RL's ``_compute_distributed_softmax``.""" + """Global softmax over a TP-sharded vocab dim. Allocates a single full-vocab output tensor (the + ``- logits_max`` subtraction) and does the rest in place on it; the input is NOT mutated (so it is + safe even when the caller passes float32 logits, where ``.float()`` is a no-op). Mirrors NeMo-RL's + ``_compute_distributed_softmax``.""" logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) - # `- logits_max` allocates once (frees the input copy); everything after is in place. - exp_logits = (vocab_parallel_logits - logits_max).exp_() + exp_logits = (vocab_parallel_logits - logits_max).exp_() # new buffer, then in place sum_exp = exp_logits.sum(dim=-1, keepdim=True) torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=group) return exp_logits.div_(sum_exp) def _vocab_parallel_log_softmax(vocab_parallel_logits, group): - """Global log-softmax over a TP-sharded vocab dim. Mirrors NeMo-RL's - ``_compute_distributed_log_softmax``; the single ``.exp()`` temporary is summed and freed.""" + """Global log-softmax over a TP-sharded vocab dim. Uses ``torch.logsumexp`` (a fused + ``[.,.,V]->[.,.,1]`` reduction) so it never materializes a full-vocab ``exp`` temporary -- the key + memory fix at a 248K vocab, where that temp is multiple GiB (NeMo-RL's + ``_compute_distributed_log_softmax`` allocates it via an explicit ``.exp().sum()``). Allocates one + full-vocab output buffer; the input is NOT mutated.""" logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) - shifted = vocab_parallel_logits - logits_max - sum_exp = shifted.exp().sum(dim=-1, keepdim=True) - torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=group) - return shifted.sub_(sum_exp.log_()) + shifted = vocab_parallel_logits - logits_max # new buffer (input untouched); values now <= 0 + # exp() of the *local* logsumexp recovers this shard's sum-of-exp, which reduces across TP. + local_sum_exp = torch.logsumexp(shifted, dim=-1, keepdim=True).exp_() + torch.distributed.all_reduce(local_sum_exp, op=torch.distributed.ReduceOp.SUM, group=group) + return shifted.sub_(local_sum_exp.log_()) # in place on the owned buffer -> log-softmax class _VocabParallelSoftCrossEntropy(torch.autograd.Function): @@ -115,17 +120,63 @@ def forward(ctx, student_vp_logits, teacher_vp_logits, tp_group): @staticmethod def backward(ctx, grad_output): student_probs, target_probs = ctx.saved_tensors - # d(H(p, q))/d(student_logit_v) = softmax(student)_v - softmax(teacher)_v - grad_student = (student_probs - target_probs) * grad_output.unsqueeze(-1) + # d(H(p, q))/d(student_logit_v) = softmax(student)_v - softmax(teacher)_v. Done in place on + # student_probs (a private buffer from forward, used nowhere else) to avoid allocating extra + # full-vocab temporaries during backward -- at a 248K vocab each would be multiple GiB. + grad_student = student_probs.sub_(target_probs).mul_(grad_output.unsqueeze(-1)) return grad_student.to(ctx.input_dtype), None, None +def _masked_mean_or_global(per_token_loss, mask, global_normalization_factor): + """Masked mean of a ``[batch, seq]`` per-token loss, using either the local token count or a + provided global valid-token count as the denominator.""" + if global_normalization_factor is not None: + return (per_token_loss * mask).sum() / global_normalization_factor.clamp(min=1.0) + return masked_mean(per_token_loss, mask) + + +def _chunked_masked_loss(per_token_fn, sliceable_args, mask, chunk_size, global_normalization_factor): + """Masked-mean loss computed in ``chunk_size`` slices along the sequence dim. + + ``per_token_fn(*sliced_args)`` returns the ``[batch, chunk]`` per-token loss for a slice. Each + chunk is gradient-checkpointed, so the (large, full-vocab) softmax tensors are recomputed in + backward instead of all being held at once -- this bounds peak memory to a *single* chunk's vocab + tensors. Required for large-vocab models (Qwen3.5's 248K vocab) where the full-sequence softmax + OOMs even after the lean kernel. The result is numerically identical to the un-chunked masked mean + (same global denominator); only the activation memory differs. + """ + import torch.utils.checkpoint as checkpoint + + seq_len = sliceable_args[0].shape[1] + masked_sum = None + for start in range(0, seq_len, chunk_size): + end = min(start + chunk_size, seq_len) + sliced = [a[:, start:end] for a in sliceable_args] + chunk_mask = mask[:, start:end] + + def _chunk_masked_sum(*args, _m=chunk_mask): + return (per_token_fn(*args) * _m).sum() + + chunk_sum = checkpoint.checkpoint(_chunk_masked_sum, *sliced, use_reentrant=False) + masked_sum = chunk_sum if masked_sum is None else masked_sum + chunk_sum + + if masked_sum is None: # empty sequence + masked_sum = mask.sum() * 0.0 + denom = ( + global_normalization_factor.clamp(min=1.0) + if global_normalization_factor is not None + else mask.sum().clamp(min=1.0) + ) + return masked_sum / denom + + def draft_soft_ce( student_logits: torch.Tensor, teacher_logits: torch.Tensor, mask: torch.Tensor, global_normalization_factor: Optional[torch.Tensor] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + chunk_size: Optional[int] = None, ) -> torch.Tensor: """Masked-mean soft cross-entropy between draft (student) and policy (teacher). @@ -138,22 +189,137 @@ def draft_soft_ce( When ``None``, uses the local masked mean. vocab_parallel_group: TP group when logits are vocab-sharded; ``None`` for full-vocab logits (single device / FSDP). + chunk_size: If set, compute the loss in sequence-chunks of this size with gradient + checkpointing to bound peak memory (large-vocab models). ``None`` computes the whole + sequence at once. The result is identical either way. Returns: Scalar loss. """ - if vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1: - per_token_loss = _VocabParallelSoftCrossEntropy.apply( - student_logits, teacher_logits, vocab_parallel_group + use_vp = vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1 + + def per_token(student, teacher): + if use_vp: + return _VocabParallelSoftCrossEntropy.apply(student, teacher, vocab_parallel_group) + teacher_probs = F.softmax(teacher.float(), dim=-1) + student_log_probs = F.log_softmax(student.float(), dim=-1) + return -(teacher_probs * student_log_probs).sum(dim=-1) + + if chunk_size is not None and student_logits.shape[1] > chunk_size: + return _chunked_masked_loss( + per_token, (student_logits, teacher_logits), mask, chunk_size, global_normalization_factor ) - else: - teacher_probs = F.softmax(teacher_logits.float(), dim=-1) - student_log_probs = F.log_softmax(student_logits.float(), dim=-1) - per_token_loss = -(teacher_probs * student_log_probs).sum(dim=-1) + return _masked_mean_or_global(per_token(student_logits, teacher_logits), mask, global_normalization_factor) - if global_normalization_factor is not None: - return (per_token_loss * mask).sum() / global_normalization_factor.clamp(min=1.0) - return masked_mean(per_token_loss, mask) + +class _VocabParallelTopkSoftCE(torch.autograd.Function): + """Top-k approximation of the soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` + that never materializes a full-vocab softmax. + + Only the teacher's top-k tokens are distilled, with teacher and student renormalized over that set. + Memory is ``O(seq * k)`` instead of ``O(seq * vocab)`` -- it fits at a 248K vocab *and* avoids the + allocator fragmentation that the full-vocab / chunked paths cause. + + Parallelism: the vocab is sharded over the tensor-parallel ``group``. We take **each rank's local + top-k** of its shard and reconcile the softmax normalizers and the per-token cross-entropy across + the group with ``all_reduce`` (MAX for the stable-softmax shift, SUM for the denominators and the + loss). No ``all_gather`` / global-index gather is needed because student and teacher share the same + vocab partition, so a rank's local top-k indices address both. This works for ANY tensor-parallel + size and is correct across nodes (the collectives run over the configured ``group``); ``group=None`` + / world-size 1 is the single-device path. The distilled set is the union of the per-rank top-k + (<= ``k * tp_size`` tokens) -- a benign, richer signal at larger TP, not a correctness issue. + + Backward is manual (scatter the ``softmax(student) - softmax(teacher)`` gradient back to this rank's + own top-k columns), so no gradient flows through the collectives. + """ + + @staticmethod + def forward(ctx, student_logits, teacher_logits, k, group, roll_shift): + ws = torch.distributed.get_world_size(group) if group is not None else 1 + + def ar(t, op): + if ws > 1: + torch.distributed.all_reduce(t, op=op, group=group) + return t + + k_eff = min(int(k), teacher_logits.shape[-1]) + # Teacher's local top-k (teacher is detached). When roll_shift != 0, teacher_logits is the + # UN-rolled policy logits: position t's draft target is the policy distribution at t+roll_shift, + # so we top-k the policy logits once (no full-vocab copy) and roll only the small [B, S, k] + # result -- avoiding a ~[S, vocab] rolled-teacher copy. The wrapped boundary positions are + # zeroed by the caller's shifted loss mask. + t_vals, t_idx = teacher_logits.topk(k_eff, dim=-1) + if roll_shift: + t_vals = torch.roll(t_vals, shifts=-int(roll_shift), dims=1) + t_idx = torch.roll(t_idx, shifts=-int(roll_shift), dims=1) + t_vals = t_vals.float() + # student at position t, teacher's top-k indices for position t (already rolled). + s_vals = student_logits.gather(-1, t_idx).float() + + # Stable-softmax shift = global max over the union (across the TP group). + t_max = ar(t_vals.max(dim=-1, keepdim=True).values.clone(), torch.distributed.ReduceOp.MAX) + s_max = ar(s_vals.max(dim=-1, keepdim=True).values.clone(), torch.distributed.ReduceOp.MAX) + + # Teacher probs over the union (denominator summed across the group). + t_exp = (t_vals - t_max).exp() + t_denom = ar(t_exp.sum(dim=-1, keepdim=True), torch.distributed.ReduceOp.SUM) + t_p = t_exp / t_denom + + # Student probs / log-probs over the union (denominator summed across the group). + s_exp = (s_vals - s_max).exp() + s_denom = ar(s_exp.sum(dim=-1, keepdim=True), torch.distributed.ReduceOp.SUM) + q_s = s_exp / s_denom + s_logp = (s_vals - s_max) - s_denom.log() + + # Per-rank partial CE summed across the group -> soft CE over the global union. + per_token_loss = ar(-(t_p * s_logp).sum(dim=-1), torch.distributed.ReduceOp.SUM) + + ctx.save_for_backward(q_s, t_p, t_idx) + ctx.vocab_size = student_logits.shape[-1] + ctx.input_dtype = student_logits.dtype + return per_token_loss + + @staticmethod + def backward(ctx, grad_output): + q_s, t_p, t_idx = ctx.saved_tensors + # d(H(p,q))/d(student_logit_v) = softmax(student)_v - softmax(teacher)_v, over the union; zero + # elsewhere. Scatter the k per-token grads back to this rank's own vocab columns. + grad_topk = (q_s - t_p) * grad_output.unsqueeze(-1) + grad_student = torch.zeros( + *t_idx.shape[:-1], ctx.vocab_size, dtype=grad_topk.dtype, device=grad_topk.device + ) + grad_student.scatter_(-1, t_idx, grad_topk) + return grad_student.to(ctx.input_dtype), None, None, None, None + + +def draft_soft_ce_topk( + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + mask: torch.Tensor, + k: int, + global_normalization_factor: Optional[torch.Tensor] = None, + vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, + roll_shift: int = 0, +) -> torch.Tensor: + """Masked-mean **top-k** soft cross-entropy between draft (student) and policy (teacher). + + Distills only the teacher's top-``k`` tokens (renormalized over that set) -- ``O(seq*k)`` memory, + no full-vocab softmax, so it fits + avoids fragmentation at large vocab. See + :class:`_VocabParallelTopkSoftCE` for the (parallelism-scalable) implementation. Approximate: drops + the teacher tail, which is benign for a draft head (acceptance is governed by the top tokens). + + Args mirror :func:`draft_soft_ce`, plus: + k: number of top teacher tokens to distill. + roll_shift: if non-zero, ``teacher_logits`` is the *un-rolled* policy logits and the draft + target for position ``t`` is the policy distribution at ``t + roll_shift`` (= the MTP-depth + roll). Top-k is taken on the un-rolled logits and only the small ``[B, S, k]`` result is + rolled, avoiding a full ``[S, vocab]`` rolled-teacher copy. ``0`` means ``teacher_logits`` + is already aligned/rolled (the caller's ``build_teacher_logits`` path). + """ + per_token_loss = _VocabParallelTopkSoftCE.apply( + student_logits, teacher_logits.detach(), k, vocab_parallel_group, roll_shift + ) + return _masked_mean_or_global(per_token_loss, mask, global_normalization_factor) def _onehot_vp_logits( @@ -186,6 +352,7 @@ def draft_hard_ce( global_normalization_factor: Optional[torch.Tensor] = None, vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, vocab_start_index: Optional[int] = None, + chunk_size: Optional[int] = None, ) -> torch.Tensor: """Masked-mean hard cross-entropy of the draft head against next-token labels. @@ -200,19 +367,23 @@ def draft_hard_ce( global_normalization_factor: see :func:`draft_soft_ce`. vocab_parallel_group: TP group when logits are vocab-sharded. vocab_start_index: This rank's vocab offset (required when TP-sharded). + chunk_size: see :func:`draft_soft_ce`. Sequence-chunked + checkpointed when set. """ - if vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1: + use_vp = vocab_parallel_group is not None and torch.distributed.get_world_size(vocab_parallel_group) > 1 + if use_vp: assert vocab_start_index is not None, "vocab_start_index is required for vocab-parallel hard CE" - # Hard CE == soft CE with a one-hot teacher; reuse the distributed soft-CE - # path so the global (TP) softmax / gradient logic lives in one place. - teacher_onehot = _onehot_vp_logits(labels, student_logits, vocab_start_index) - per_token_loss = _VocabParallelSoftCrossEntropy.apply( - student_logits, teacher_onehot, vocab_parallel_group - ) - else: - log_probs = F.log_softmax(student_logits.float(), dim=-1) - per_token_loss = -log_probs.gather(dim=-1, index=labels.unsqueeze(-1).long()).squeeze(-1) - if global_normalization_factor is not None: - return (per_token_loss * mask).sum() / global_normalization_factor.clamp(min=1.0) - return masked_mean(per_token_loss, mask) + def per_token(student, lbl): + if use_vp: + # Hard CE == soft CE with a one-hot teacher; reuse the distributed soft-CE path so the + # global (TP) softmax / gradient logic lives in one place. + teacher_onehot = _onehot_vp_logits(lbl, student, vocab_start_index) + return _VocabParallelSoftCrossEntropy.apply(student, teacher_onehot, vocab_parallel_group) + log_probs = F.log_softmax(student.float(), dim=-1) + return -log_probs.gather(dim=-1, index=lbl.unsqueeze(-1).long()).squeeze(-1) + + if chunk_size is not None and student_logits.shape[1] > chunk_size: + return _chunked_masked_loss( + per_token, (student_logits, labels), mask, chunk_size, global_normalization_factor + ) + return _masked_mean_or_global(per_token(student_logits, labels), mask, global_normalization_factor) diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index a542fe1acf..ba7d11460d 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -229,6 +229,21 @@ class MegatronConfig(BaseConfig): mtp_loss_scaling_factor: Optional[float] = None """Deprecated alias for ``mtp_loss_weight`` (kept for back-compat). If set, it overrides ``mtp_loss_weight``.""" + mtp_loss_chunk_size: Optional[int] = 1024 + """Sequence-chunk size for the decoupled MTP/draft loss. The draft loss materializes full + ``[batch, seq, vocab]`` softmax tensors; at large vocab (e.g. Qwen3.5's 248K) computing the whole + response at once OOMs. Chunking the loss over the sequence (with gradient checkpointing) bounds + peak memory to one chunk's vocab tensors -- numerically identical, only activation memory differs. + ``None`` disables chunking (whole sequence at once). Lower it if the draft loss still OOMs. + Ignored when ``mtp_loss_topk`` is set (top-k never materializes the full vocab, so no chunking).""" + mtp_loss_topk: Optional[int] = None + """If set, use a **top-k** approximation of the soft-CE draft loss: distill only the teacher's top-k + tokens (renormalized over that set) instead of the full vocabulary. Memory is ``O(seq*k)`` instead + of ``O(seq*vocab)``, which both fits and avoids allocator fragmentation at large vocab (e.g. + Qwen3.5's 248K) -- where even the chunked full-vocab loss OOMs/fragments. Scales to any + tensor-parallel size (incl. cross-node): the per-rank top-k is reconciled across the TP group with + all-reduce. Only applies to ``mtp_loss_type="soft_ce"``. ``None`` uses the exact full-vocab loss. + Typical: 64-128 (acceptance is governed by the top tokens, so the dropped tail is benign).""" def __post_init__(self): # Back-compat: the old `mtp_loss_scaling_factor` knob now maps onto the explicit loss weight. diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 24efcb8987..f6574c18f2 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -768,6 +768,7 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: "UV_CACHE_DIR", "UV_LINK_MODE", "UV_PYTHON", + "UV_OFFLINE", "MTP_DEBUG", "PYTORCH_CUDA_ALLOC_CONF", ): diff --git a/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py b/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py index 30404cadf3..4c238497f7 100644 --- a/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py +++ b/tests/backends/skyrl_train/inference_engines/test_inference_engine_client.py @@ -354,61 +354,8 @@ async def generate(self, input_batch): assert out["stop_reasons"][i] == "stop" -# ------------------------------------------- -# tests for InferenceEngineClient.get_spec_decode_metrics -# -------------------------------------------- - - -def _make_spec_client(per_engine_stats): - """Build a client whose engines report the given per-engine spec-decode stats.""" - - class MockEngine: - def __init__(self, stats): - self._stats = stats - - def dp_size(self): - return 1 - - async def get_spec_decode_metrics(self): - return self._stats - - cfg = _make_min_cfg() - return InferenceEngineClient( - engines=[MockEngine(s) for s in per_engine_stats], - tokenizer=object(), - model_path=cfg.trainer.policy.model.path, - lora_cfg=cfg.trainer.policy.model.lora, - inference_engine_cfg=cfg.generator.inference_engine, - ) - - -@pytest.mark.asyncio -async def test_get_spec_decode_metrics_sums_across_engines(): - client = _make_spec_client( - [ - {"num_drafts": 3, "num_draft_tokens": 10, "num_accepted_tokens": 6}, - {"num_drafts": 1, "num_draft_tokens": 4, "num_accepted_tokens": 2}, - ] - ) - totals = await client.get_spec_decode_metrics() - assert totals == {"num_drafts": 4, "num_draft_tokens": 14, "num_accepted_tokens": 8} - - -@pytest.mark.asyncio -async def test_get_spec_decode_metrics_returns_none_when_no_engine_reports(): - # Speculative decoding disabled / unsupported on every engine -> None (not a zero-dict). - client = _make_spec_client([None, None]) - assert await client.get_spec_decode_metrics() is None - - -@pytest.mark.asyncio -async def test_get_spec_decode_metrics_ignores_non_reporting_engines(): - # A mix of reporting and non-reporting engines still aggregates the reporters only. - client = _make_spec_client( - [None, {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5}] - ) - totals = await client.get_spec_decode_metrics() - assert totals == {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5} +# NOTE: tests for InferenceEngineClient.get_spec_decode_metrics (MTP draft) live in +# tests/backends/skyrl_train/mtp/test_spec_decode_client.py # ----------------------------- diff --git a/tests/backends/skyrl_train/mtp/test_soft_ce.py b/tests/backends/skyrl_train/mtp/test_soft_ce.py index adac63fdb7..8a7a70929e 100644 --- a/tests/backends/skyrl_train/mtp/test_soft_ce.py +++ b/tests/backends/skyrl_train/mtp/test_soft_ce.py @@ -59,6 +59,126 @@ def test_vocab_parallel_soft_ce_preserves_input_dtype(monkeypatch): assert student.grad.dtype == torch.bfloat16 +def test_draft_soft_ce_chunked_matches_unchunked(): + # Sequence-chunking (+ gradient checkpointing) must be numerically identical to the whole-sequence + # loss in both value and gradient -- it only bounds activation memory. + torch.manual_seed(0) + student = torch.randn(2, 9, 7) + teacher = torch.randn(2, 9, 7) + mask = torch.ones(2, 9) + mask[0, 5:] = 0 # partial mask to exercise the masked denominator across chunk boundaries + + s1 = student.clone().requires_grad_(True) + loss_full = draft_soft_ce(s1, teacher, mask) + loss_full.backward() + + s2 = student.clone().requires_grad_(True) + loss_chunked = draft_soft_ce(s2, teacher, mask, chunk_size=4) # 9 -> chunks of 4,4,1 + loss_chunked.backward() + + assert torch.allclose(loss_full, loss_chunked, atol=1e-6) + assert torch.allclose(s1.grad, s2.grad, atol=1e-6) + + +def test_draft_hard_ce_chunked_matches_unchunked(): + torch.manual_seed(1) + student = torch.randn(2, 9, 7) + labels = torch.randint(0, 7, (2, 9)) + mask = torch.ones(2, 9) + mask[1, 6:] = 0 + + s1 = student.clone().requires_grad_(True) + loss_full = draft_hard_ce(s1, labels, mask) + loss_full.backward() + + s2 = student.clone().requires_grad_(True) + loss_chunked = draft_hard_ce(s2, labels, mask, chunk_size=4) + loss_chunked.backward() + + assert torch.allclose(loss_full, loss_chunked, atol=1e-6) + assert torch.allclose(s1.grad, s2.grad, atol=1e-6) + + +def test_draft_soft_ce_chunk_size_larger_than_seq_is_noop(): + # chunk_size >= seq_len must behave exactly like the un-chunked path. + torch.manual_seed(2) + student = torch.randn(1, 4, 5, requires_grad=True) + teacher = torch.randn(1, 4, 5) + mask = torch.ones(1, 4) + assert torch.allclose( + draft_soft_ce(student, teacher, mask, chunk_size=999), + draft_soft_ce(student, teacher, mask), + atol=1e-6, + ) + + +def test_draft_soft_ce_topk_matches_reference(): + # Single-device (TP=1) top-k soft CE must equal a reference: distill the teacher's top-k tokens, + # renormalized over that set. Checks both value and gradient (the custom backward scatters + # softmax(student)-softmax(teacher) to the top-k columns). + from skyrl.backends.skyrl_train.mtp.soft_ce import draft_soft_ce_topk + + torch.manual_seed(0) + student = torch.randn(2, 5, 11) + teacher = torch.randn(2, 5, 11) + mask = torch.ones(2, 5) + mask[0, 3:] = 0 + k = 4 + + s1 = student.clone().requires_grad_(True) + got = draft_soft_ce_topk(s1, teacher, mask, k=k) + got.backward() + + # Reference: gather student at teacher's top-k, softmax/CE over the k set. + s2 = student.clone().requires_grad_(True) + t_vals, t_idx = teacher.topk(k, dim=-1) + s_vals = s2.gather(-1, t_idx) + t_p = F.softmax(t_vals, dim=-1) + ref_per_token = -(t_p * F.log_softmax(s_vals, dim=-1)).sum(-1) + ref = (ref_per_token * mask).sum() / mask.sum() + ref.backward() + + assert torch.allclose(got, ref, atol=1e-6) + assert torch.allclose(s1.grad, s2.grad, atol=1e-6) + + +def test_draft_soft_ce_topk_roll_shift_matches_prerolled(): + # roll_shift (top-k on the un-rolled policy logits, then roll the small [B,S,k] result) must equal + # pre-rolling the full teacher then top-k with roll_shift=0 -- in both value and gradient. This is + # the memory optimization that avoids the full [S, vocab] rolled-teacher copy. + from skyrl.backends.skyrl_train.mtp.soft_ce import draft_soft_ce_topk + + torch.manual_seed(0) + student = torch.randn(2, 6, 11) + teacher = torch.randn(2, 6, 11) + mask = torch.ones(2, 6) + shift, k = 2, 4 + + s1 = student.clone().requires_grad_(True) + got = draft_soft_ce_topk(s1, teacher, mask, k=k, roll_shift=shift) + got.backward() + + s2 = student.clone().requires_grad_(True) + pre_rolled = torch.roll(teacher, shifts=-shift, dims=1) + ref = draft_soft_ce_topk(s2, pre_rolled, mask, k=k, roll_shift=0) + ref.backward() + + assert torch.allclose(got, ref, atol=1e-6) + assert torch.allclose(s1.grad, s2.grad, atol=1e-6) + + +def test_draft_soft_ce_topk_memory_is_topk_sized(): + # The forward must not materialize a full-vocab tensor: gradient is nonzero only at the k columns. + from skyrl.backends.skyrl_train.mtp.soft_ce import draft_soft_ce_topk + + torch.manual_seed(1) + student = torch.randn(1, 3, 20, requires_grad=True) + teacher = torch.randn(1, 3, 20) + draft_soft_ce_topk(student, teacher, torch.ones(1, 3), k=5).backward() + # exactly k nonzero grad columns per token. + assert int((student.grad != 0).sum(-1).max()) <= 5 + + def test_soft_ce_matches_reference(): torch.manual_seed(0) student = torch.randn(2, 5, 7, requires_grad=True) diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_client.py b/tests/backends/skyrl_train/mtp/test_spec_decode_client.py new file mode 100644 index 0000000000..85269cd9d9 --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_client.py @@ -0,0 +1,89 @@ +"""CPU unit tests for InferenceEngineClient.get_spec_decode_metrics aggregation (MTP draft). + +The client sums the per-engine cumulative spec-decode counters; the generator turns the per-step +delta into an acceptance rate (see test_spec_decode_metrics.py for that piece). + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_spec_decode_client.py +""" + +import pytest + +from skyrl.backends.skyrl_train.inference_engines.inference_engine_client import ( + InferenceEngineClient, +) +from skyrl.train.config import ( + GeneratorConfig, + InferenceEngineConfig, + ModelConfig, + PolicyConfig, + SkyRLTrainConfig, + TrainerConfig, +) + + +def _make_min_cfg(): + return SkyRLTrainConfig( + trainer=TrainerConfig( + policy=PolicyConfig(model=ModelConfig(path="dummy-model")), + ), + generator=GeneratorConfig( + inference_engine=InferenceEngineConfig( + backend="vllm", + enable_http_endpoint=False, + http_endpoint_host="127.0.0.1", + http_endpoint_port=0, + ), + ), + ) + + +def _make_spec_client(per_engine_stats): + """Build a client whose engines report the given per-engine spec-decode stats.""" + + class MockEngine: + def __init__(self, stats): + self._stats = stats + + def dp_size(self): + return 1 + + async def get_spec_decode_metrics(self): + return self._stats + + cfg = _make_min_cfg() + return InferenceEngineClient( + engines=[MockEngine(s) for s in per_engine_stats], + tokenizer=object(), + model_path=cfg.trainer.policy.model.path, + lora_cfg=cfg.trainer.policy.model.lora, + inference_engine_cfg=cfg.generator.inference_engine, + ) + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_sums_across_engines(): + client = _make_spec_client( + [ + {"num_drafts": 3, "num_draft_tokens": 10, "num_accepted_tokens": 6}, + {"num_drafts": 1, "num_draft_tokens": 4, "num_accepted_tokens": 2}, + ] + ) + totals = await client.get_spec_decode_metrics() + assert totals == {"num_drafts": 4, "num_draft_tokens": 14, "num_accepted_tokens": 8} + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_returns_none_when_no_engine_reports(): + # Speculative decoding disabled / unsupported on every engine -> None (not a zero-dict). + client = _make_spec_client([None, None]) + assert await client.get_spec_decode_metrics() is None + + +@pytest.mark.asyncio +async def test_get_spec_decode_metrics_ignores_non_reporting_engines(): + # A mix of reporting and non-reporting engines still aggregates the reporters only. + client = _make_spec_client( + [None, {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5}] + ) + totals = await client.get_spec_decode_metrics() + assert totals == {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5} From 2440d2864bf8d917a91a9756a9dc340d896079c1 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 10 Jun 2026 03:23:16 +0000 Subject: [PATCH 10/28] fix oom --- .../new_inference_worker_wrap.py | 10 +++++ .../inference_servers/spec_decode_utils.py | 42 +++++++++++++++++++ .../inference_servers/vllm_worker.py | 5 +++ skyrl/backends/skyrl_train/mtp/soft_ce.py | 26 ++++++++++-- .../megatron/megatron_model_wrapper.py | 2 +- .../skyrl_train/workers/worker_utils.py | 21 ++++++++-- skyrl/train/utils/utils.py | 6 +++ 7 files changed, 104 insertions(+), 8 deletions(-) create mode 100644 skyrl/backends/skyrl_train/inference_servers/spec_decode_utils.py diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index 9fe64f4c40..5a8ef65277 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -169,6 +169,16 @@ def update_weights_ipc(self, update_info: dict) -> None: with set_current_vllm_config(self.vllm_config), torch.device(self.device): if self._skyrl_is_checkpoint_format: model.load_weights(weights=weights) + # vLLM's load only updates the main model; the spec-decode (MTP/Eagle) + # drafter is a separate module and must be reloaded from the same + # checkpoint-format weights (see spec_decode_utils). No-op without spec + # decoding. Only valid for checkpoint-format names (the drafter's + # load_weights expects them); the per-param copy branch is skipped. + from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( + _reload_spec_decode_drafter, + ) + + _reload_spec_decode_drafter(self.model_runner, weights) else: for name, weight in weights: param = model.get_parameter(name) diff --git a/skyrl/backends/skyrl_train/inference_servers/spec_decode_utils.py b/skyrl/backends/skyrl_train/inference_servers/spec_decode_utils.py new file mode 100644 index 0000000000..7cd9c03b6c --- /dev/null +++ b/skyrl/backends/skyrl_train/inference_servers/spec_decode_utils.py @@ -0,0 +1,42 @@ +"""Spec-decode (MTP / Eagle) drafter weight-sync helper. + +vLLM's weight reload (``GPUModelRunner.reload_weights`` and direct +``model_runner.model.load_weights``) only updates the **main** model. For models +with native MTP heads (Qwen3.5, DeepSeek-V3, GLM-4.x, ...) the speculative-decoding +**drafter is a SEPARATE module** (``model_runner.drafter.model``, e.g. +``Qwen3_5MTP``) that the main-model load never touches. In colocated training the +inference engine is also slept with ``level=2`` (which discards weights), so after +wake the drafter is left uninitialized/garbage and MTP speculative decoding drafts +with a broken head -> ~0 draft-acceptance from the first generate. + +This helper re-loads the drafter from the same synced weights right after the main +model load. The drafter's ``load_weights`` filters to the names it consumes (e.g. +Qwen3.5: ``mtp.*`` plus ``embed_tokens``/``lm_head``) and ignores the rest, so the +full weight list is safe to pass. No-op when speculative decoding is disabled (no +drafter) or the drafter has no loadable model (e.g. ngram). +""" + +from typing import Iterable, List, Tuple + +import torch + + +def _reload_spec_decode_drafter(model_runner, weight_list: List[Tuple[str, torch.Tensor]]) -> bool: + """Re-load the spec-decode drafter model from ``weight_list`` (HF-named tensors). + + Args: + model_runner: the vLLM ``GPUModelRunner`` (``self.model_runner`` on the worker). + weight_list: the list of ``(name, tensor)`` pairs already received for the + main-model sync. A list (not a one-shot iterator) so it can be re-iterated. + + Returns: + True if a drafter model was reloaded, False if there was nothing to reload. + """ + drafter = getattr(model_runner, "drafter", None) + drafter_model = getattr(drafter, "model", None) + if drafter_model is None or not hasattr(drafter_model, "load_weights"): + # No spec decoding, or a drafter without a weight-loadable model (e.g. ngram). + return False + weights: Iterable[Tuple[str, torch.Tensor]] = iter(weight_list) + drafter_model.load_weights(weights) + return True diff --git a/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py b/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py index 463caf2477..b621bd615c 100644 --- a/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py +++ b/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py @@ -84,6 +84,10 @@ def load_weights(self, request: bytes) -> None: from vllm.config import set_current_vllm_config + from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( + _reload_spec_decode_drafter, + ) + # Unpickle request to restore the original object type assert isinstance(request, bytes), f"Expected bytes, got {type(request).__name__}" request = pickle.loads(request) @@ -94,6 +98,7 @@ def load_weights(self, request: bytes) -> None: with torch.device(self.device), set_current_vllm_config(self.vllm_config): self.model_runner.reload_weights(weights_iterator=iter(weight_list)) + _reload_spec_decode_drafter(self.model_runner, weight_list) for weight in weight_list: del weight diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index ffb70c04b2..1bfc087156 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -51,14 +51,32 @@ def build_teacher_logits( def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.Tensor: """Roll a ``[batch, seq]`` loss mask to align with an MTP teacher/label. - Positions that wrap around the sequence end have no valid target, so they - are zeroed. This mirrors how Megatron's ``roll_tensor`` zeros the wrapped - boundary inside ``process_mtp_loss``. + MTP depth ``k`` supervises position ``t`` against a target read at ``t + k + 1`` + (the policy distribution / label for ``seq[t + k + 2]``). For that supervision to + be valid, BOTH the *source* position ``t`` (where the draft head produced its + logits) and the *target* position ``t + k + 1`` (where the teacher/label is read) + must be real tokens. We therefore require ``mask[t] AND mask[t + shift]``: + + * ``rolled = roll(mask, -shift)`` is the target-side validity ``mask[t + shift]``, + with the wrapped tail zeroed (no valid target past the sequence end). + * multiplying by the original ``mask`` re-applies the source-side validity. + + The source-side ``AND`` is what makes this correct under **left padding** (and any + padding layout): rolling a left-padded mask left turns the last pad slot into a + ``1`` (its target is the first real token), which would otherwise pull a de-padded + *zero-logit* (→ uniform) pad position into the loss and inflate it by up to + ``log(V)`` per leaked position. Re-ANDing the source mask drops those positions, so + the loss no longer depends on the de-pad zero-fill happening to be masked out. + + Note: each de-padded row holds a single sequence, so the one wrapped boundary is + handled by zeroing ``rolled``'s tail. If multiple sequences are ever packed into a + single row, per-sequence boundaries would need ``roll_tensor``-style cu_seqlens + handling here. """ shift = mtp_layer_number + 1 rolled = torch.roll(mask, shifts=-shift, dims=1) rolled[:, -shift:] = 0 - return rolled + return mask * rolled def _vocab_parallel_softmax(vocab_parallel_logits, group): diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 9926364202..c5da00b341 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -395,7 +395,7 @@ def loss_func(logits, data): chunk_size=mtp_loss_chunk_size, ) ) - draft_loss = torch.stack(per_layer_losses).mean() + draft_loss = torch.stack(per_layer_losses).mean() # Drop the dict's reference, this microbatch's autograd graph still holds the # tensor for its own backward, after which it is freed instead of lingering. del data["mtp_student_logits"] diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index cecaccc17e..382c22962f 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -6,12 +6,23 @@ from skyrl.train.dataset.replay_buffer import Experience +# Metrics that end in `_loss` but are plain per-token MEANS, not pre-scaled minibatch sums. +# The `sum_loss_metrics` convention sums every `_loss` key because the *policy* losses are +# pre-scaled (by num_microbatches * dp_size) so that summing recovers the correct minibatch +# loss. The decoupled MTP/draft loss is NOT pre-scaled -- it is a masked per-token mean, like +# KL/entropy (which dodge the sum only because they are named `policy_kl`/`policy_entropy`). +# Summing it would multiply the reported value by the microbatch (and DP) count -- e.g. a true +# ~0.5 nats reads as ~44. Keep these averaged regardless of `sum_loss_metrics`. +MEAN_LOSS_METRICS = frozenset({"mtp_loss", "draft_loss"}) + + def reduce_metrics(metrics: Dict[str, List[float]], sum_loss_metrics: bool = False) -> Dict[str, float]: """Reduce scalar metrics from a list of entries per key with the appropriate reduction. Default reduction is mean. Metrics ending in `_min` or `_max` use min/max respectively. - If sum_loss_metrics is True, metrics ending in `_loss` are summed instead of averaged. + If sum_loss_metrics is True, metrics ending in `_loss` are summed instead of averaged + (except those in `MEAN_LOSS_METRICS`, which are always averaged). This should be used if the scaling is already done at the advantage level. See `apply_loss_reduction_to_advantages_minibatch` for more details. @@ -30,7 +41,7 @@ def reduce_metrics(metrics: Dict[str, List[float]], sum_loss_metrics: bool = Fal reduced_metrics[k] = max(v) elif k.endswith("_min"): reduced_metrics[k] = min(v) - elif sum_loss_metrics and k.endswith("_loss"): + elif sum_loss_metrics and k.endswith("_loss") and k not in MEAN_LOSS_METRICS: reduced_metrics[k] = sum(v) else: reduced_metrics[k] = sum(v) / len(v) @@ -56,7 +67,11 @@ def all_reduce_metrics( """ min_metrics = {k: v for k, v in metrics.items() if k.endswith("_min")} max_metrics = {k: v for k, v in metrics.items() if k.endswith("_max")} - sum_metrics = {k: v for k, v in metrics.items() if sum_loss_metrics and k.endswith("_loss")} + sum_metrics = { + k: v + for k, v in metrics.items() + if sum_loss_metrics and k.endswith("_loss") and k not in MEAN_LOSS_METRICS + } mean_metrics = { k: v for k, v in metrics.items() if k not in min_metrics and k not in max_metrics and k not in sum_metrics } diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index f6574c18f2..ec7a3225be 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -771,6 +771,12 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: "UV_OFFLINE", "MTP_DEBUG", "PYTORCH_CUDA_ALLOC_CONF", + # Debug/trace knobs — forwarded so they reach the worker actors, not just the driver. + "CUDA_LAUNCH_BLOCKING", + "PYTHONFAULTHANDLER", + "TORCH_SHOW_CPP_STACKTRACES", + "TORCH_USE_CUDA_DSA", + "NCCL_DEBUG", ): if value := os.environ.get(var_name): logger.info(f"Exporting `{var_name}` to ray runtime env: {value}") From ca33d92621815d9067b848e941d0dd2741f9b244 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 10 Jun 2026 21:28:40 +0000 Subject: [PATCH 11/28] fix roll mismatch in Hard CE, k > 1 mtp support --- .../inference_engine_client.py | 15 +++-- .../vllm/spec_decode_metrics.py | 54 +++++++++++++++-- skyrl/backends/skyrl_train/mtp/adapter.py | 13 ++-- .../skyrl_train/mtp/draft_loss_wrapper.py | 6 +- .../skyrl_train/mtp/hidden_capture.py | 37 ++++++++++-- skyrl/backends/skyrl_train/mtp/soft_ce.py | 40 ++++++++++--- .../megatron/megatron_model_wrapper.py | 59 +++++++++++++++---- .../workers/megatron/megatron_worker.py | 20 +++++++ .../skyrl_train/workers/worker_utils.py | 5 +- skyrl/train/config/config.py | 29 +++++---- skyrl/train/trainer.py | 6 +- skyrl/train/utils/utils.py | 22 +++++-- 12 files changed, 245 insertions(+), 61 deletions(-) diff --git a/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py b/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py index 828fb32f24..cbd9bb1000 100644 --- a/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py +++ b/skyrl/backends/skyrl_train/inference_engines/inference_engine_client.py @@ -383,20 +383,25 @@ async def resume_generation(self) -> None: """ await self._run_on_all_engines("resume_generation") - async def get_spec_decode_metrics(self) -> Optional[Dict[str, int]]: + async def get_spec_decode_metrics(self) -> Optional[Dict[str, Any]]: """Sum cumulative speculative-decoding counters across all engines. - Returns None if no engine reports spec-decode stats (e.g. speculative decoding disabled). + Scalar counters add; per-position accepted-count lists add elementwise (see + ``merge_spec_decode_counters``). Returns None if no engine reports spec-decode stats + (e.g. speculative decoding disabled). """ + from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( + merge_spec_decode_counters, + ) + per_engine = await self._run_on_all_engines("get_spec_decode_metrics") - totals: Dict[str, int] = {} + totals: Dict[str, Any] = {} any_reported = False for stats in per_engine: if not stats: continue any_reported = True - for key, value in stats.items(): - totals[key] = totals.get(key, 0) + int(value) + merge_spec_decode_counters(totals, stats) return totals if any_reported else None # ---------------------------- diff --git a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py index 637d56c1de..1295f87d61 100644 --- a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py +++ b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py @@ -3,12 +3,25 @@ # vLLM's async engine (the one SkyRL uses for generation + weight sync) does not expose # ``get_metrics()``, so we attach a tiny custom stat logger that runs in the AsyncLLM frontend # process (= the SkyRL engine actor) and accumulates the per-iteration spec-decode counters -# (``num_draft_tokens`` / ``num_accepted_tokens``). The engine reads these cumulative counters and -# the generator turns them into a per-step acceptance rate (accepted / drafted). +# (``num_draft_tokens`` / ``num_accepted_tokens`` / ``num_accepted_tokens_per_pos``). The engine +# reads these cumulative counters and the generator turns them into per-step acceptance rates: +# one overall rate (accepted / drafted) plus one rate per draft position (how often the k-th +# speculated token was accepted), so multi-token drafting (num_speculative_tokens > 1) shows the +# per-depth acceptance decay. The number of per-position metrics tracks whatever depth vLLM +# reports, so it grows/shrinks with the configured draft depth automatically. from __future__ import annotations -from typing import Optional +from typing import Any, Optional + + +def _add_per_pos(into: list, add) -> None: + """Elementwise-add a per-position count list into ``into``, growing ``into`` as needed.""" + add = add or [] + if len(add) > len(into): + into.extend([0] * (len(add) - len(into))) + for i, n in enumerate(add): + into[i] += int(n) def make_spec_decode_stat_logger_class(): @@ -27,6 +40,10 @@ def __init__(self, vllm_config, engine_index: int = 0): self.num_drafts = 0 self.num_draft_tokens = 0 self.num_accepted_tokens = 0 + # Accepted count per draft position (index 0 = first speculated token). vLLM reports a + # list of length num_speculative_tokens each iteration; we grow on demand instead of + # pre-sizing so the logger needs no knowledge of the configured depth. + self.num_accepted_tokens_per_pos: list = [] def record(self, scheduler_stats=None, iteration_stats=None, mm_cache_stats=None, engine_idx: int = 0): stats = getattr(scheduler_stats, "spec_decoding_stats", None) if scheduler_stats is not None else None @@ -34,6 +51,7 @@ def record(self, scheduler_stats=None, iteration_stats=None, mm_cache_stats=None self.num_drafts += stats.num_drafts self.num_draft_tokens += stats.num_draft_tokens self.num_accepted_tokens += stats.num_accepted_tokens + _add_per_pos(self.num_accepted_tokens_per_pos, getattr(stats, "num_accepted_tokens_per_pos", None)) def log_engine_initialized(self): pass @@ -45,13 +63,30 @@ def sum_spec_decode_loggers(loggers) -> Optional[dict]: """Sum cumulative counters across a list of ``SpecDecodeStatLogger`` instances.""" if not loggers: return None + per_pos: list = [] + for lg in loggers: + _add_per_pos(per_pos, getattr(lg, "num_accepted_tokens_per_pos", None)) return { "num_drafts": sum(int(lg.num_drafts) for lg in loggers), "num_draft_tokens": sum(int(lg.num_draft_tokens) for lg in loggers), "num_accepted_tokens": sum(int(lg.num_accepted_tokens) for lg in loggers), + "num_accepted_tokens_per_pos": per_pos, } +def merge_spec_decode_counters(totals: dict, stats: dict) -> None: + """Merge one engine's counter dict into ``totals`` in place (cross-engine aggregation). + + Scalar counters add; per-position lists add elementwise (padding to the longest list, so + engines configured with different draft depths still merge correctly). + """ + for key, value in stats.items(): + if isinstance(value, list): + _add_per_pos(totals.setdefault(key, []), value) + else: + totals[key] = totals.get(key, 0) + int(value) + + def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> tuple[dict, Optional[dict]]: """Turn cumulative spec-decode counters into per-step (delta) metrics. @@ -67,13 +102,24 @@ def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> if not cumulative: return {}, prev prev = prev or {} + drafts = cumulative.get("num_drafts", 0) - prev.get("num_drafts", 0) drafted = cumulative.get("num_draft_tokens", 0) - prev.get("num_draft_tokens", 0) accepted = cumulative.get("num_accepted_tokens", 0) - prev.get("num_accepted_tokens", 0) - metrics = { + metrics: dict[str, Any] = { "vllm/draft_num_draft_tokens": drafted, "vllm/draft_num_accepted_tokens": accepted, } if drafted > 0: # Acceptance rate = accepted draft tokens / total drafted tokens this step. metrics["vllm/draft_acceptance_rate"] = accepted / drafted + # Per-position rates: fraction of draft rounds this step whose k-th speculated token was + # accepted (same definition vLLM uses in its own per-position logging). Position keys are + # 1-based: pos_1 = first drafted token. Acceptance halts at the first rejection, so the rates + # are non-increasing in k — the decay shows how much each extra draft position actually pays. + per_pos = cumulative.get("num_accepted_tokens_per_pos") or [] + prev_per_pos = prev.get("num_accepted_tokens_per_pos") or [] + if drafts > 0: + for i, n in enumerate(per_pos): + prev_n = int(prev_per_pos[i]) if i < len(prev_per_pos) else 0 + metrics[f"vllm/draft_acceptance_rate_pos_{i + 1}"] = (int(n) - prev_n) / drafts return metrics, cumulative diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py index af33168a47..085197376a 100644 --- a/skyrl/backends/skyrl_train/mtp/adapter.py +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -57,10 +57,15 @@ def project_mtp_hidden_to_logits(hidden_states_per_layer, model, detach_output_w output_weight = None if getattr(model, "share_embeddings_and_output_weights", False): output_weight = model.shared_embedding_or_output_weight() - if detach_output_weight: - # Optionally isolate the shared embedding/output-layer from the draft - # gradient too (default keeps NeMo's behaviour: only the trunk hidden - # states are detached). + if detach_output_weight: + # Isolate the output projection from the draft gradient: detach the shared/tied weight, + # or — for untied models — pass the output layer's own weight explicitly as a detached + # tensor (ColumnParallelLinear uses a provided ``weight`` verbatim). Mirrors slime's + # MTP-RL megatron patch. Combined with the capture's detached re-embedding, the draft + # loss then trains only the MTP-head parameters. + if output_weight is None: + output_weight = getattr(model.output_layer, "weight", None) + if output_weight is not None: output_weight = output_weight.detach() logits_per_layer = [] diff --git a/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py b/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py index d44c7894fd..9e519914f1 100644 --- a/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py +++ b/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py @@ -64,7 +64,11 @@ def combine_policy_and_draft_loss( """ draft_losses = [] for layer_idx, student_logits in enumerate(student_logits_per_layer): - layer_mask = shift_mask_for_mtp(mask, layer_idx) + # Hard CE's target is the *token* ``seq[t+k+2]`` — one position past the soft-CE teacher + # (a *distribution* read at ``t+k+1``) — so its validity mask needs one extra shift, or the + # last in-range position of every row trains against a zeroed/pad label. + mask_shift = layer_idx + 1 if cfg.loss_type == "hard_ce" else layer_idx + layer_mask = shift_mask_for_mtp(mask, mask_shift) if cfg.loss_type == "hard_ce": assert hard_labels is not None, "hard_ce requires hard_labels" layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index 43f85fb0cb..60538ae229 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -74,7 +74,9 @@ def _mtp_layer_offset(mtp_block) -> int: the common single-stage case; we resolve it via megatron-core so the replay split below is also correct under PP/VPP. Falls back to 0 if the helper is unavailable (older megatron-core).""" try: - from megatron.core.transformer.multi_token_prediction import get_mtp_layer_offset + from megatron.core.transformer.multi_token_prediction import ( + get_mtp_layer_offset, + ) return int(get_mtp_layer_offset(mtp_block.config, getattr(mtp_block, "vp_stage", None))) except Exception: @@ -90,13 +92,20 @@ class MTPHiddenCapture: output layer. """ - def __init__(self, model, detach_trunk: bool = True): + def __init__(self, model, detach_trunk: bool = True, detach_shared_embedding: bool = False): # Resolve the module that owns the MTP block: the unwrapped model itself for plain # GPT/Mamba models, or the nested ``.language_model`` for VL wrappers (e.g. Qwen3.5-VL). # capture.model is also what the caller projects through (output_layer / shared weight), # so it must be the same module that holds the heads. self.model = _resolve_mtp_host(_unwrap_model(model)) self.detach_trunk = detach_trunk + # The MTP block re-embeds the rolled input ids through the model's shared embedding + # (``embedding=self.embedding`` in its kwargs) — a second gradient path into the shared + # embedding weight, separate from the output projection. When the caller wants the shared + # weights fully isolated from the draft loss (``mtp_detach_shared_output``), the replay must + # detach this path too, else tied-embedding models still train the lm_head through the + # lookup (slime's MTP-RL patch detaches the same two paths). + self.detach_shared_embedding = detach_shared_embedding self._args = None self._kwargs = None self._handles: list = [] @@ -148,9 +157,25 @@ def compute_student_hidden_states(self) -> Optional[List]: import torch kwargs = dict(self._kwargs) - hidden = kwargs.get("hidden_states") - if hidden is not None and self.detach_trunk: + if self.detach_trunk: + hidden = kwargs.get("hidden_states") + # The detach below only patches the *keyword* argument. Megatron currently calls + # ``self.mtp(...)`` with kwargs only (GPTModel + MambaModel); if a future version passes + # hidden_states positionally, the detach would silently no-op and the draft gradient + # would couple back into the policy trunk — fail loudly instead. + assert hidden is not None, ( + "MTP capture: 'hidden_states' was not passed to the MTP block as a keyword argument " + "(megatron-core call convention changed?); cannot detach the trunk for decoupled " + "draft training." + ) kwargs["hidden_states"] = hidden.detach() + if self.detach_shared_embedding and kwargs.get("embedding") is not None: + orig_embedding = kwargs["embedding"] + + def _detached_embedding(*emb_args, **emb_kwargs): + return orig_embedding(*emb_args, **emb_kwargs).detach() + + kwargs["embedding"] = _detached_embedding mtp = self.model.mtp # MultiTokenPredictionBlock concatenates, along dim 0: @@ -165,11 +190,11 @@ def compute_student_hidden_states(self) -> Optional[List]: @contextmanager -def maybe_capture_mtp_hidden(model, enabled: bool, detach_trunk: bool = True): +def maybe_capture_mtp_hidden(model, enabled: bool, detach_trunk: bool = True, detach_shared_embedding: bool = False): """Context manager returning an ``MTPHiddenCapture`` when ``enabled``, else ``None``.""" if not enabled: yield None return - capture = MTPHiddenCapture(model, detach_trunk=detach_trunk) + capture = MTPHiddenCapture(model, detach_trunk=detach_trunk, detach_shared_embedding=detach_shared_embedding) with capture.capture(): yield capture diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index 1bfc087156..376e0fbb2b 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -79,6 +79,30 @@ def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.T return mask * rolled +def unpadded_vocab_shard_width( + true_vocab_size: Optional[int], + vocab_shard_width: int, + tp_rank: int, +) -> int: + """Width of this rank's vocab shard excluding Megatron's padded tail. + + When the global vocab is padded to divide across TP (``should_pad_vocab``), the padding occupies + the tail of the last rank(s)' shards (the vocab is partitioned contiguously). Padded weight rows + are never trained (zero-filled by the HF conversion), so their logits must not enter the draft + loss: in the teacher softmax they leak ~uniform probability mass, and in the teacher top-k they + can steal slots from real tokens. Callers slice ``logits[..., :width]`` — a view, so autograd + zero-fills the padded tail's gradient and per-rank widths may differ (the loss collectives only + reduce ``[..., 1]``/``[...]``-shaped tensors, never the vocab dim). + + Returns ``vocab_shard_width`` unchanged when ``true_vocab_size`` is ``None`` (unknown) or the + shard holds no padding. + """ + if true_vocab_size is None: + return vocab_shard_width + start = tp_rank * vocab_shard_width + return max(0, min(vocab_shard_width, true_vocab_size - start)) + + def _vocab_parallel_softmax(vocab_parallel_logits, group): """Global softmax over a TP-sharded vocab dim. Allocates a single full-vocab output tensor (the ``- logits_max`` subtraction) and does the rest in place on it; the input is NOT mutated (so it is @@ -301,13 +325,13 @@ def ar(t, op): def backward(ctx, grad_output): q_s, t_p, t_idx = ctx.saved_tensors # d(H(p,q))/d(student_logit_v) = softmax(student)_v - softmax(teacher)_v, over the union; zero - # elsewhere. Scatter the k per-token grads back to this rank's own vocab columns. + # elsewhere. Scatter the k per-token grads back to this rank's own vocab columns. The full-vocab + # buffer is allocated directly in the input dtype (the only fp32->input cast is the tiny [.., k] + # grad); an fp32 buffer here would double the largest transient of the whole loss. grad_topk = (q_s - t_p) * grad_output.unsqueeze(-1) - grad_student = torch.zeros( - *t_idx.shape[:-1], ctx.vocab_size, dtype=grad_topk.dtype, device=grad_topk.device - ) - grad_student.scatter_(-1, t_idx, grad_topk) - return grad_student.to(ctx.input_dtype), None, None, None, None + grad_student = torch.zeros(*t_idx.shape[:-1], ctx.vocab_size, dtype=ctx.input_dtype, device=grad_topk.device) + grad_student.scatter_(-1, t_idx, grad_topk.to(ctx.input_dtype)) + return grad_student, None, None, None, None def draft_soft_ce_topk( @@ -401,7 +425,5 @@ def per_token(student, lbl): return -log_probs.gather(dim=-1, index=lbl.unsqueeze(-1).long()).squeeze(-1) if chunk_size is not None and student_logits.shape[1] > chunk_size: - return _chunked_masked_loss( - per_token, (student_logits, labels), mask, chunk_size, global_normalization_factor - ) + return _chunked_masked_loss(per_token, (student_logits, labels), mask, chunk_size, global_normalization_factor) return _masked_mean_or_global(per_token(student_logits, labels), mask, global_normalization_factor) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index c5da00b341..73d958eec3 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -22,14 +22,6 @@ from_parallel_logits_to_logprobs, vocab_parallel_entropy, ) -from skyrl.backends.skyrl_train.utils.ppo_utils import ( - PolicyLossRegistry, - compute_approx_kl, -) -from skyrl.backends.skyrl_train.utils.replay_utils import ( - setup_per_microbatch_replay_backward, - setup_per_microbatch_replay_forward, -) from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits from skyrl.backends.skyrl_train.mtp.hidden_capture import maybe_capture_mtp_hidden from skyrl.backends.skyrl_train.mtp.soft_ce import ( @@ -38,6 +30,15 @@ draft_soft_ce, draft_soft_ce_topk, shift_mask_for_mtp, + unpadded_vocab_shard_width, +) +from skyrl.backends.skyrl_train.utils.ppo_utils import ( + PolicyLossRegistry, + compute_approx_kl, +) +from skyrl.backends.skyrl_train.utils.replay_utils import ( + setup_per_microbatch_replay_backward, + setup_per_microbatch_replay_forward, ) from skyrl.backends.skyrl_train.utils.torch_utils import ( build_mtp_next_token_labels, @@ -257,10 +258,20 @@ def forward_backward_mini_batch( mtp_loss_type = getattr(mcfg, "mtp_loss_type", "soft_ce") mtp_loss_weight = float(getattr(mcfg, "mtp_loss_weight", 0.1)) mtp_detach_trunk = bool(getattr(mcfg, "mtp_detach_trunk", True)) - mtp_detach_shared_output = bool(getattr(mcfg, "mtp_detach_shared_output", False)) + mtp_detach_shared_output = bool(getattr(mcfg, "mtp_detach_shared_output", True)) mtp_loss_chunk_size = getattr(mcfg, "mtp_loss_chunk_size", 1024) mtp_loss_topk = getattr(mcfg, "mtp_loss_topk", None) + if mtp_enabled: + # The decoupled draft training records the MTP block's inputs with a forward pre-hook + # and replays the block afterwards; module hooks do not fire inside CUDA-graph replay, + # so the capture would silently see nothing and the draft loss would vanish. + assert ( + not getattr(model_config, "enable_cuda_graph", False) + and not getattr(model_config, "external_cuda_graph", False) + and getattr(model_config, "cuda_graph_impl", "none") in (None, "none") + ), "MTP draft training uses forward hooks and cannot be combined with Megatron CUDA graphs" + if os.environ.get("MTP_DEBUG"): from megatron.core.utils import unwrap_model as _uw @@ -350,11 +361,29 @@ def loss_func(logits, data): # Undo the in-place temperature scaling so the teacher is the true policy # distribution (student logits are produced unscaled). teacher_src = logits if temperature == 1.0 else logits * temperature + # Megatron may pad the global vocab to divide across TP (should_pad_vocab); padded + # weight rows are never trained, so their logits would leak probability mass into + # the teacher softmax and could steal teacher top-k slots. Slice this rank's shard + # to its true width (a view; autograd zero-fills the tail's grad). model_config here + # is the bridge provider, whose vocab_size is the true (unpadded) HF vocab; the + # vocab_start_index passed to hard CE keeps the PADDED stride (global column offset). + true_shard_width = unpadded_vocab_shard_width( + getattr(model_config, "vocab_size", None), vocab_size_tp, tp_rank + ) + if true_shard_width != vocab_size_tp: + teacher_src = teacher_src[..., :true_shard_width] hard_labels = build_mtp_next_token_labels(sequences) if mtp_loss_type == "hard_ce" else None per_layer_losses = [] for layer_idx, student_logits in enumerate(student_logits_list): - layer_mask = shift_mask_for_mtp(draft_mask, layer_idx) + if true_shard_width != vocab_size_tp: + student_logits = student_logits[..., :true_shard_width] + # The hard label for position t at depth k is the *token* seq[t+k+2] — one + # position past the soft-CE teacher (a *distribution* read at t+k+1) — so its + # validity mask needs one extra shift. Without it, the last in-range position of + # every row trains against a zeroed/pad label. + mask_shift = layer_idx + 1 if mtp_loss_type == "hard_ce" else layer_idx + layer_mask = shift_mask_for_mtp(draft_mask, mask_shift) if mtp_loss_type == "hard_ce": layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) per_layer_losses.append( @@ -610,7 +639,15 @@ def depad(tensor): # main logits stay coupled to the trunk; MTPLossAutoScaler never runs). student_hidden = None student_model = None - with maybe_capture_mtp_hidden(model, mtp_enabled, detach_trunk=mtp_detach_trunk) as capture: + with maybe_capture_mtp_hidden( + model, + mtp_enabled, + detach_trunk=mtp_detach_trunk, + # Full shared-weight isolation also requires detaching the MTP block's re-embedding + # of the rolled input ids (the second gradient path into the tied embedding/lm_head, + # next to the output projection handled in project_mtp_hidden_to_logits). + detach_shared_embedding=mtp_detach_shared_output, + ) as capture: outputs = model( new_sequences, new_position_ids, diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 8eda134d7b..0b1088f48e 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -418,6 +418,26 @@ def init_configs( provider.mtp_num_layers = None elif megatron_config.mtp_num_layers is not None: provider.mtp_num_layers = megatron_config.mtp_num_layers or None + # When MTP training is requested (trainer.mtp.enabled), the model must actually resolve to + # >= 1 head — otherwise the draft loss would silently no-op (the wrapper keys off + # model_config.mtp_num_layers). _apply_mtp_config no longer forces a head count, so a model + # whose checkpoint ships no MTP layers lands here with mtp_num_layers=None; users who want + # to train fresh heads from scratch can still force-build them via + # policy.megatron_config.mtp_num_layers (note vLLM's `method: mtp` would also need the HF + # config to declare them). + mtp_cfg = getattr(self.cfg, "mtp", None) + if ( + enable_mtp + and mtp_cfg is not None + and getattr(mtp_cfg, "enabled", False) + and not getattr(provider, "mtp_num_layers", None) + ): + raise ValueError( + "trainer.mtp.enabled=true but the model resolved to 0 MTP heads " + "(the checkpoint's HF config declares none and policy.megatron_config.mtp_num_layers " + "is unset). Use an MTP-capable checkpoint, or set " + "policy.megatron_config.mtp_num_layers to force-build fresh heads." + ) if getattr(provider, "mtp_num_layers", None): # The native MTP heads are still *built* (so the megatron-bridge weight mappings # round-trip them to HF / vLLM), but SkyRL trains them with a decoupled, explicit diff --git a/skyrl/backends/skyrl_train/workers/worker_utils.py b/skyrl/backends/skyrl_train/workers/worker_utils.py index 382c22962f..1be9446dfb 100644 --- a/skyrl/backends/skyrl_train/workers/worker_utils.py +++ b/skyrl/backends/skyrl_train/workers/worker_utils.py @@ -5,7 +5,6 @@ from skyrl.backends.skyrl_train.training_batch import TrainingInputBatch from skyrl.train.dataset.replay_buffer import Experience - # Metrics that end in `_loss` but are plain per-token MEANS, not pre-scaled minibatch sums. # The `sum_loss_metrics` convention sums every `_loss` key because the *policy* losses are # pre-scaled (by num_microbatches * dp_size) so that summing recovers the correct minibatch @@ -68,9 +67,7 @@ def all_reduce_metrics( min_metrics = {k: v for k, v in metrics.items() if k.endswith("_min")} max_metrics = {k: v for k, v in metrics.items() if k.endswith("_max")} sum_metrics = { - k: v - for k, v in metrics.items() - if sum_loss_metrics and k.endswith("_loss") and k not in MEAN_LOSS_METRICS + k: v for k, v in metrics.items() if sum_loss_metrics and k.endswith("_loss") and k not in MEAN_LOSS_METRICS } mean_metrics = { k: v for k, v in metrics.items() if k not in min_metrics and k not in max_metrics and k not in sum_metrics diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index ba7d11460d..366fe6cda0 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -223,9 +223,14 @@ class MegatronConfig(BaseConfig): """If True (default, matches NeMo-RL), detach the trunk hidden states feeding the MTP head so only the MTP-head parameters receive the draft gradient. If False, let the draft gradient flow back into the policy backbone (closer to Megatron's native coupled MTP).""" - mtp_detach_shared_output: bool = False - """If True, also detach the shared embedding/output-layer weight from the draft gradient. - Default keeps the shared head trainable by the draft loss (only the trunk hidden is detached).""" + mtp_detach_shared_output: bool = True + """If True (default), fully isolate the shared embedding/output weights from the draft gradient: + the draft-loss projection uses a detached output weight AND the MTP block's re-embedding of the + rolled input ids is detached, so only the MTP-head parameters (enorm/hnorm/eh_proj/layer/final + norm) train. This matters for tied-embedding models (e.g. Qwen3.5), where a trainable shared + head means the draft loss directly nudges the policy's own logits every step. Matches slime's + MTP-RL training (only ``.mtp.`` params receive gradients). Set False for NeMo-RL's behaviour + (shared embedding/head also trained by the draft loss).""" mtp_loss_scaling_factor: Optional[float] = None """Deprecated alias for ``mtp_loss_weight`` (kept for back-compat). If set, it overrides ``mtp_loss_weight``.""" @@ -684,11 +689,12 @@ class MTPConfig(BaseConfig): """High-level Multi-Token Prediction (MTP) control. This is the single user-facing knob for MTP. When ``enabled``, ``validate_cfg`` propagates it to - the training side (``policy.megatron_config.mtp_*`` — build and train ``num_speculative_tokens`` - native MTP heads with a decoupled draft loss) and the inference side - (``generator.inference_engine.speculative_config`` — vLLM MTP speculative decoding with the same - number of draft tokens). The trained heads are kept in sync with the policy via weight sync, so - the draft tracks the updated policy. + the training side (``policy.megatron_config.mtp_*`` — train the model's native MTP heads with a + decoupled draft loss; the head count comes from the checkpoint's own HF config, overridable via + ``policy.megatron_config.mtp_num_layers``) and the inference side + (``generator.inference_engine.speculative_config`` — vLLM MTP speculative decoding with + ``num_speculative_tokens`` draft tokens per step). The trained heads are kept in sync with the + policy via weight sync, so the draft tracks the updated policy. The trunk hidden states and embeddings feeding the heads are always detached (decoupled) — that is intrinsic to draft/Eagle training, not a tunable, so it is not exposed here.""" @@ -696,8 +702,11 @@ class MTPConfig(BaseConfig): enabled: bool = False """Whether to train MTP draft heads and use them for speculative decoding.""" num_speculative_tokens: int = 1 - """Number of future tokens to predict = number of MTP heads / draft depth. Must be <= the number - of MTP layers the model ships (e.g. GLM-4.7-Flash / DeepSeek-V3 ``num_nextn_predict_layers``).""" + """Draft depth: how many tokens vLLM speculates per decoding step. Decoupled from the number of + trained MTP heads — with a single-head checkpoint (Qwen3.5/Qwen3-Next/DeepSeek-V3 all ship one), + depths > 1 reuse that head autoregressively at draft time (vLLM never indexes beyond head 0). + Expect per-position acceptance to decay with depth, since the head is trained at depth 1 on true + trunk hidden states but consumes its own outputs at deeper draft positions.""" loss_type: str = "soft_ce" """``"soft_ce"`` (distill against the policy's own next-token distribution) or ``"hard_ce"``.""" loss_weight: float = 0.1 diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index b60c315303..d7913fa53e 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -845,8 +845,10 @@ async def _record_spec_decode_metrics(self) -> None: """Record the vLLM speculative-decoding (MTP draft) acceptance rate for this rollout step. Reads the engines' cumulative draft/accept counters and logs the per-step delta as - ``vllm/draft_acceptance_rate`` (+ raw counts). No-op when speculative decoding is disabled or - the backend doesn't expose the stats. Best-effort: never fails the training step. + ``vllm/draft_acceptance_rate`` (+ raw counts), plus one ``vllm/draft_acceptance_rate_pos_k`` + per draft position when num_speculative_tokens > 1 (per-depth acceptance decay). No-op when + speculative decoding is disabled or the backend doesn't expose the stats. Best-effort: + never fails the training step. """ from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( acceptance_rate_metrics, diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index ec7a3225be..b30e50c667 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -221,9 +221,14 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): def _apply_mtp_config(cfg: SkyRLTrainConfig): """Propagate the high-level ``trainer.mtp`` knob to the training + inference configs. - When MTP is enabled, build/train ``num_speculative_tokens`` native MTP heads (Megatron) with the - decoupled draft loss, and enable vLLM MTP speculative decoding with the same draft depth. When - disabled, force the MTP heads off so they are neither built nor run. + When MTP is enabled, train the model's native MTP heads (Megatron builds as many as the + checkpoint ships — typically one; override via ``policy.megatron_config.mtp_num_layers``) with + the decoupled draft loss, and enable vLLM MTP speculative decoding with + ``num_speculative_tokens`` draft tokens per step. The draft depth is deliberately decoupled + from the trained head count: vLLM drafts depth > 1 by reusing the (single) MTP head + autoregressively, so e.g. ``num_speculative_tokens=3`` with a 1-head checkpoint is valid and + does NOT build extra (randomly initialized) heads on the training side. When disabled, force + the MTP heads off so they are neither built nor run. """ mtp = getattr(cfg.trainer, "mtp", None) if mtp is None: @@ -236,8 +241,15 @@ def _apply_mtp_config(cfg: SkyRLTrainConfig): return assert mtp.num_speculative_tokens >= 1, "trainer.mtp.num_speculative_tokens must be >= 1 when enabled" - # Training side: build + train this many native MTP heads with the decoupled draft loss. - mcfg.mtp_num_layers = mtp.num_speculative_tokens + if mcfg.mtp_num_layers == 0: + raise ValueError( + "trainer.mtp.enabled=true but trainer.policy.megatron_config.mtp_num_layers=0 " + "(explicit force-disable). Remove the mtp_num_layers override or disable trainer.mtp." + ) + # Training side: train the checkpoint's native MTP heads with the decoupled draft loss. + # mcfg.mtp_num_layers is intentionally left untouched (None => the megatron-bridge infers the + # head count from the model's own HF config); MegatronWorker fails loudly if the model resolves + # to zero heads while MTP is enabled. mcfg.mtp_loss_type = mtp.loss_type mcfg.mtp_loss_weight = mtp.loss_weight From 84931418c4bea5a267f752164b4fd24db6192fe7 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 10 Jun 2026 23:37:36 +0000 Subject: [PATCH 12/28] add test cases and running scripts --- .ray_cache_bust | 1 + ...run_megatron_dapo_qwen3.5_2b_specdecode.sh | 39 ++- .../megatron/run_megatron_dapo_qwen3.5_9b.sh | 156 +++++++++++ .../megatron/test_mtp_packed_vs_unpacked.py | 254 ++++++++++++++++++ .../megatron/test_mtp_replay_vs_native.py | 215 +++++++++++++++ .../megatron/test_mtp_weight_roundtrip.py | 242 +++++++++++++++++ .../test_spec_decode_drafter_reload.py | 67 +++++ .../skyrl_train/mtp/test_hidden_capture.py | 98 ++++++- .../skyrl_train/mtp/test_mtp_config.py | 41 ++- .../backends/skyrl_train/mtp/test_soft_ce.py | 122 ++++++++- .../mtp/test_spec_decode_client.py | 4 +- .../mtp/test_spec_decode_metrics.py | 130 +++++++-- .../skyrl_train/workers/test_worker_utils.py | 13 + 13 files changed, 1347 insertions(+), 35 deletions(-) create mode 100644 .ray_cache_bust create mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh create mode 100644 tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py create mode 100644 tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_replay_vs_native.py create mode 100644 tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py create mode 100644 tests/backends/skyrl_train/inference_servers/test_spec_decode_drafter_reload.py diff --git a/.ray_cache_bust b/.ray_cache_bust new file mode 100644 index 0000000000..4a78c7e881 --- /dev/null +++ b/.ray_cache_bust @@ -0,0 +1 @@ +1780815368 diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh index a86dec2a7c..80bb179000 100644 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh @@ -55,6 +55,9 @@ TOP_P=1.0 EVAL_TOP_P=0.7 CLIP_RATIO_C=10.0 MAX_PROMPT_LENGTH=$((1024 * 2)) +# IDENTICAL to the no-spec run (8192) for an apples-to-apples comparison. The MTP draft loss adds a +# full [seq, vocab] softmax tensor on top of the main lm-head gradient, but this fits at TP=2 with the +# default allocator now that the O(num_microbatches) `mtp_student_logits` accumulation is fixed. MAX_RESPONSE_LENGTH=$((1024 * 8)) # repro run parameters @@ -69,9 +72,13 @@ ENFORCE_EAGER=true # cuda graphs can cause some instability LR=1e-6 # megatron config -- Qwen3.5-2B is a dense model, so no expert parallelism. -# TP=2 (not 1): with 8K-token responses, TP=2 auto-enables sequence parallelism, -# which shards activations/vocab-logits across the TP group. TP=1 keeps full -# activations on each GPU and OOMs in the backward grad-sync at gpu_mem_util=0.5. +# TP=2 (matching the no-spec run): the decoupled MTP draft loss materializes a full +# [batch, seq, vocab] float32 softmax tensor on top of the main lm-head gradient at Qwen3.5's 248K +# vocab, but this fits at TP=2 with the default allocator now that the O(num_microbatches) +# `mtp_student_logits` accumulation is fixed. NOTE: raising TP does NOT meaningfully shrink the +# training footprint (it's activation/optimizer-bound, not vocab-bound), and TP=8 also triggered a +# CUDA-IPC weight-sync failure (pidfd_getfd EPERM) when broadcasting the 8-way-sharded policy to the +# TP=1 vLLM engines. So TP=2 is both the apples-to-apples match and the known-good weight-sync config. # TP=2, PP=1, CP=1 => DP=4. MEGATRON_TP=2 MEGATRON_PP=1 @@ -86,11 +93,25 @@ TIS_TYPE=token # Multi-Token Prediction (MTP) speculative decoding. -# Qwen3.5-2B ships 1 native MTP head (`mtp_num_hidden_layers: 1`), so draft depth = 1. +# Qwen3.5-2B ships 1 native MTP head (`mtp_num_hidden_layers: 1`); training always trains the +# checkpoint's heads. NUM_SPECULATIVE_TOKENS is the vLLM *draft depth* only — values > 1 reuse the +# single head autoregressively at draft time (more speedup, but per-position acceptance decays with +# depth since the head never trains on its own outputs). Try 2-3 for extra acceleration. MTP_ENABLED=true MTP_NUM_SPECULATIVE_TOKENS=1 MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) MTP_LOSS_WEIGHT=0.1 +# NOTE: trainer.policy.megatron_config.mtp_detach_shared_output now defaults to true: the draft loss +# trains ONLY the MTP-head params; the tied embedding/lm_head is detached (output projection AND the +# MTP block's re-embedding), so the draft gradient no longer nudges the policy's own logits. Set it +# to false for the old NeMo-RL behaviour (shared head also trained by the draft loss). +# Top-k draft loss: distill only the teacher's top-k tokens instead of the full 248K vocab, keeping +# draft-loss memory at O(seq*k) vs O(seq*vocab). This is now an optional throughput knob, NOT a +# memory requirement: the full-vocab soft-CE fits comfortably at TP=2 since the real OOM cause (the +# O(num_microbatches) accumulation of `mtp_student_logits`) was fixed by freeing each microbatch's +# student logits right after its draft backward. No PYTORCH_CUDA_ALLOC_CONF tuning is needed. +# null => exact full-vocab soft-CE. Set to e.g. 64 to use the memory-light top-k approximation. +MTP_LOSS_TOPK=null # Qwen3.5 flags @@ -99,6 +120,15 @@ ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/ DISTRIBUTED_EXECUTOR_BACKEND="mp" export _SKYRL_USE_NEW_INFERENCE=0 export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 +# NOTE: do NOT set PYTORCH_CUDA_ALLOC_CONF here (neither expandable_segments:True nor max_split_size_mb). +# - expandable_segments:True makes PyTorch allocate via virtual-memory segments incompatible with the +# legacy CUDA-IPC handle used by colocated weight sync -> it falls back to pidfd_getfd, which this +# cluster's ptrace_scope blocks (pidfd_getfd: Operation not permitted). +# - max_split_size_mb over-reserves PyTorch memory and starves NCCL's external cudaMalloc at grad-sync +# ("Failed to CUDA calloc ... bytes" in reduce_scatter). +# Neither is needed: the draft-loss OOM was an O(num_microbatches) accumulation of `mtp_student_logits`, +# now fixed by freeing each microbatch's student logits right after its draft backward. The full-vocab +# soft-CE fits at TP=2 with the default allocator. uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ data.train_data="['$TRAIN_FILE']" \ @@ -165,6 +195,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ trainer.mtp.loss_type=$MTP_LOSS_TYPE \ trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ + trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ trainer.logger="$LOGGER" \ trainer.project_name="qwen3_5_dapo_sd" \ trainer.run_name="sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh new file mode 100644 index 0000000000..446da8a885 --- /dev/null +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh @@ -0,0 +1,156 @@ +set -x + +# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO with Megatron. +# Runs on 1 node of 8xH100s (80GB each). +# +# NOTE: verify the exact HF repo id for the 9B model before running +# (e.g. `hf download Qwen/Qwen3.5-9B` / check https://huggingface.co/Qwen). +# +# Prepare data onto the fast local disk first: +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# Then launch: +# bash examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh + +MODEL_NAME="Qwen/Qwen3.5-9B" +# Use the fast, non-persistent local disk for data (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +# 9B is ~4.5x the 2B: a single full-weight copy is ~18GB in bf16. Use inference +# TP=2 (4 engines) so each engine's weights are halved (~9GB/GPU) and there is +# more headroom for KV cache during generation. TP=2 comm stays on NVLink. +NUM_INFERENCE_ENGINES=4 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=2 +LOGGER="wandb" # change to "console" to print to stdout + +CLIP_RATIO_LOW=0.2 +CLIP_RATIO_HIGH=0.28 +# use token mean loss reduction +LOSS_REDUCTION="token_mean" +# applies overlong filtering (but not soft overlong punishment) +APPLY_OVERLONG_FILTERING=true +# apply soft overlong punishment with custom trainer impl in main_dapo.py +OVERLONG_BUFFER_LEN=$((1024 * 4)) +OVERLONG_BUFFER_PENALTY_FACTOR=1.0 + +# other DAPO parameters +USE_KL_LOSS=false +TEMPERATURE=1.0 +TOP_P=1.0 +EVAL_TOP_P=0.7 +CLIP_RATIO_C=10.0 +MAX_PROMPT_LENGTH=$((1024 * 2)) +MAX_RESPONSE_LENGTH=$((1024 * 8)) + +# repro run parameters +TRAIN_BATCH_SIZE=128 +MINI_BATCH_SIZE=32 +N_SAMPLES_PER_PROMPT=8 +EVAL_N_SAMPLES_PER_PROMPT=16 +ENFORCE_EAGER=true # cuda graphs can cause some instability +LR=1e-6 + +# megatron config -- Qwen3.5-9B is a dense model, so no expert parallelism. +# TP=4 (up from 2 on the 2B): 9B params + Adam states + 8K-token activations +# need more sharding to fit at micro batch 1. TP>1 auto-enables sequence +# parallelism, sharding activations/vocab-logits across the TP group. +# TP=4, PP=1, CP=1 => DP=2. TP stays within the single-node NVLink domain. +MEGATRON_TP=4 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + +# optimizer offload -- OFF by default (with colocate_all the inference engine is +# offloaded during training, so the full 80GB is available for the optimizer). +# Flip to true if you hit OOM in the optimizer step / grad sync. +OPTIMIZER_OFFLOAD=false +OPTIMIZER_OFFLOAD_FRACTION=1.0 + +# TIS parameters +TIS_IMP_RATIO_CAP=2.0 +TIS_TYPE=token + +# Qwen3.5 flags +REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 +ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 +DISTRIBUTED_EXECUTOR_BACKEND="mp" +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="dual_clip" \ + trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ + trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + generator.sampling_params.temperature=$TEMPERATURE \ + generator.sampling_params.top_p=$TOP_P \ + generator.eval_sampling_params.top_p=$EVAL_TOP_P \ + generator.eval_sampling_params.temperature=$TEMPERATURE \ + generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=10 \ + trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ + trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=-1 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=$LR \ + trainer.policy.optimizer_config.num_warmup_steps=5 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.logger="$LOGGER" \ + trainer.project_name="qwen3_5_dapo" \ + trainer.run_name="dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.export_path="/mnt/local_storage/exports/dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.hf_save_interval=300 \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_path="/mnt/local_storage/ckpts/dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + $@ diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py new file mode 100644 index 0000000000..40397c35e1 --- /dev/null +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py @@ -0,0 +1,254 @@ +"""Localize Bug B: the MTP soft-CE is ~2.5 on a clean single UNPADDED sequence but +inflates to ~44 (TP=1) on a real batched, LEFT-PADDED RL batch -- far above ln(V)=12.4, +so the draft head looks *confidently wrong* on real training data. + +The existing `test_mtp_replay_vs_native` calls the model DIRECTLY +(`self.actor_module[0](sequences, position_ids, attention_mask)`), i.e. the BSHD path. +Real training defaults to `remove_microbatch_padding=True`, which routes through +`preprocess_packed_seqs` (THD packing). So the clean probe is structurally blind to the +packed path. This probe runs BOTH on the SAME content and reports soft-CE at the real +positions for each, plus replay-vs-native divergence under packing: + + - unpacked soft-CE ~= packed soft-CE -> bug is NOT the packed forward (look loss-side). + - packed soft-CE >> unpacked, replay ~= native -> Megatron's THD MTP forward is the + culprit (we feed it packed inputs the MTP path mishandles). + - packed soft-CE >> unpacked, replay != native -> our REPLAY diverges under packing. + +Run:: + uv run --isolated --extra megatron --extra dev pytest -s -vvv \ + tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py +""" + +import pytest +import ray +import torch + +from skyrl.backends.skyrl_train.workers.megatron import ( + megatron_worker as _megatron_worker_mod, +) +from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import ( + MegatronPolicyWorkerBase, +) +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.utils import init_worker_with_type + +MODEL_NAME = "Qwen/Qwen3.5-2B" + + +class _ProbeMegatronPolicyWorker(MegatronPolicyWorkerBase): + def probe_packed_vs_unpacked(self, token_ids_list) -> dict: + import torch + import megatron.core.parallel_state as mpu + from megatron.core.utils import unwrap_model + + from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( + recover_left_padding, + remove_left_padding, + ) + from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits + from skyrl.backends.skyrl_train.mtp.hidden_capture import ( + MTPHiddenCapture, + _mtp_layer_offset, + _resolve_mtp_host, + _unwrap_model, + ) + from skyrl.backends.skyrl_train.mtp.soft_ce import ( + build_teacher_logits, + draft_soft_ce, + shift_mask_for_mtp, + ) + + host = _resolve_mtp_host(_unwrap_model(self.actor_module[0])) + if getattr(host, "mtp", None) is None: + return {"error": "no mtp head built", "host": type(host).__name__} + num_layers = int(getattr(host.config, "mtp_num_layers", 1) or 1) + device = torch.cuda.current_device() + + def to_bsv(x): + if x.dim() != 3: + return x + return x if x.shape[0] != 1 or x.shape[1] == 1 else x # [1,S,V] already + + tp_grp = mpu.get_tensor_model_parallel_group() + tp_size = mpu.get_tensor_model_parallel_world_size() + + def _buggy_shift_mask(mask, k): + # Pre-fix behaviour: roll only, no source-side AND (leaks left-pad boundary). + shift = k + 1 + rolled = torch.roll(mask, shifts=-shift, dims=1) + rolled[:, -shift:] = 0 + return rolled + + # ---- soft-CE on a [B,S,V] main + per-depth student set ---- + # vp_group exercises the real vocab-parallel path when TP>1 (mirrors loss_func). + def soft_ce_over_depths(main_logits, students, mask, mask_fn, vp=False): + out = [] + for k, st in enumerate(students): + teacher = build_teacher_logits(main_logits, k, detach=True) + lm = mask_fn(mask, k) + out.append( + draft_soft_ce( + st, teacher, lm, + vocab_parallel_group=tp_grp if vp else None, + ).item() + ) + return out + + # ========================================================================= + # (1) UNPACKED baseline: run each sequence alone (no padding), BSHD path. + # ========================================================================= + # Skip the raw single-seq baseline under SP/TP>1: it bypasses remove_left_padding's + # seq-length padding to a TP multiple, so odd-length seqs trip the SP scatter. The batched + # path below uses remove_left_padding and is the one that matters for the TP>1 (158) case. + unpacked_soft_ce = [] + for ids in token_ids_list if tp_size == 1 else []: + seq = torch.tensor(ids, device=device, dtype=torch.long).unsqueeze(0) + S = seq.shape[1] + pos = torch.arange(S, device=device).unsqueeze(0) + am = torch.ones_like(seq) + cap = MTPHiddenCapture(self.actor_module[0], detach_trunk=True) + with torch.no_grad(), cap.capture(): + out = self.actor_module[0](seq, pos, am) + st_hidden = cap.compute_student_hidden_states() + main = to_bsv(out).float() + students = [s.float() for s in project_mtp_hidden_to_logits(st_hidden, host)] + unpacked_soft_ce.append( + soft_ce_over_depths(main, students, am.float(), shift_mask_for_mtp, vp=tp_size > 1)[0] + ) + + # ========================================================================= + # (2) BATCHED LEFT-PADDED, non-packed (remove_left_padding) -- the ACTUAL + # Qwen3.5 training path (GDN can't pack -> REMOVE_MICROBATCH_PADDING=false). + # ========================================================================= + B = len(token_ids_list) + L = max(len(x) for x in token_ids_list) + sequences = torch.zeros(B, L, device=device, dtype=torch.long) + attention_mask = torch.zeros(B, L, device=device, dtype=torch.long) + for i, ids in enumerate(token_ids_list): + n = len(ids) + sequences[i, L - n :] = torch.tensor(ids, device=device) # LEFT padding + attention_mask[i, L - n :] = 1 + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 0) + + am_bool = attention_mask.to(bool) + is_first = mpu.is_pipeline_first_stage(ignore_virtual=True) + new_sequences, new_attention_mask, new_position_ids = remove_left_padding( + sequences, am_bool, position_ids, pre_process=is_first + ) + is_last = mpu.is_pipeline_last_stage(ignore_virtual=True) + + def depad(t): + return recover_left_padding(t, new_attention_mask, am_bool, L, post_process=is_last) + + # capture in-forward NATIVE mtp output via hook, alongside our replay + inforward = {} + mtp = host.mtp + hh = mtp.register_forward_hook(lambda _m, _i, o: inforward.__setitem__("out", o.detach())) + cap = MTPHiddenCapture(self.actor_module[0], detach_trunk=True) + try: + with torch.no_grad(), cap.capture(): + out = self.actor_module[0](new_sequences, new_position_ids, new_attention_mask) + replay_hidden = cap.compute_student_hidden_states() + finally: + hh.remove() + + main_packed = depad(out).float() # [B,L,V] (forward_step de-pads the model output directly) + replay_students = [depad(s).float() for s in project_mtp_hidden_to_logits(replay_hidden, host)] + + total = 1 + _mtp_layer_offset(mtp) + num_layers + native_hidden = list(torch.chunk(inforward["out"], total, dim=0))[-num_layers:] + native_students = [depad(s).float() for s in project_mtp_hidden_to_logits(native_hidden, host)] + + mask_f = attention_mask.float() + vp = tp_size > 1 + replay_fixed = soft_ce_over_depths(main_packed, replay_students, mask_f, shift_mask_for_mtp, vp=vp) + replay_buggy = soft_ce_over_depths(main_packed, replay_students, mask_f, _buggy_shift_mask, vp=vp) + native_fixed = soft_ce_over_depths(main_packed, native_students, mask_f, shift_mask_for_mtp, vp=vp) + + # replay-vs-native divergence at REAL positions only (pad positions are zero-filled). + m3 = am_bool.unsqueeze(-1) + max_abs = max( + ((r - n).abs() * m3).max().item() for r, n in zip(replay_students, native_students) + ) + pad_frac = 1.0 - (attention_mask.sum().item() / (B * L)) + + return { + "host": type(host).__name__, + "mtp_num_layers": num_layers, + "tp_size": tp_size, + "batch_shape": (B, L), + "pad_fraction": round(pad_frac, 3), + "unpacked_soft_ce_per_seq": [round(x, 3) for x in unpacked_soft_ce], + "batched_soft_ce_replay_FIXED_mask": [round(x, 3) for x in replay_fixed], + "batched_soft_ce_replay_BUGGY_mask": [round(x, 3) for x in replay_buggy], + "batched_soft_ce_native_FIXED_mask": [round(x, 3) for x in native_fixed], + "max_abs_replay_minus_native": max_abs, + } + + +_ProbePolicyWorker = ray.remote(num_gpus=1)(_ProbeMegatronPolicyWorker) + + +def _cfg(): + import os + + tp = int(os.environ.get("PROBE_TP", "1")) + cfg = SkyRLTrainConfig() + cfg.trainer.policy.model.path = MODEL_NAME + cfg.trainer.strategy = "megatron" + cfg.trainer.logger = "console" + cfg.trainer.placement.colocate_all = False + cfg.trainer.placement.policy_num_gpus_per_node = tp + cfg.trainer.placement.ref_num_gpus_per_node = tp + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.context_parallel_size = 1 + cfg.trainer.mtp.enabled = True + cfg.trainer.mtp.num_speculative_tokens = 1 + cfg.trainer.mtp.loss_type = "soft_ce" + validate_cfg(cfg) + return cfg + + +@pytest.mark.megatron +def test_mtp_packed_vs_unpacked(ray_init_fixture): + import os + + tp = int(os.environ.get("PROBE_TP", "1")) + cfg = _cfg() + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + # Coherent passages of VERY different lengths -> heavy left padding in the batch (stresses the + # left-pad mask leak: the shorter the seq, the larger its pad fraction). + long_txt = ( + "The mitochondria is the powerhouse of the cell. Photosynthesis converts sunlight, water, " + "and carbon dioxide into glucose and oxygen. In 1969, Apollo 11 landed humans on the Moon. " + "Water boils at one hundred degrees Celsius at sea level. The capital of France is Paris." + ) + short_txt = "The quick brown fox jumps over the lazy dog near the river bank at dawn." + tiny_txt = "Paris is the capital of France." + ids_long = tok(long_txt, add_special_tokens=False)["input_ids"][:96] + ids_short = tok(short_txt, add_special_tokens=False)["input_ids"][:40] + ids_tiny = tok(tiny_txt, add_special_tokens=False)["input_ids"][:10] + token_ids_list = [ids_long, ids_short, ids_tiny] + + _orig = _megatron_worker_mod.PolicyWorker + _megatron_worker_mod.PolicyWorker = _ProbePolicyWorker + try: + policy = init_worker_with_type( + "policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp, cfg=cfg + ) + res = ray.get( + policy.async_run_ray_method("pass_through", "probe_packed_vs_unpacked", token_ids_list) + )[0] + finally: + _megatron_worker_mod.PolicyWorker = _orig + + print("\n===== MTP packed vs unpacked =====") + for k, v in res.items(): + print(f"{k:36s}: {v}") + print("==================================") + assert "error" not in res, res diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_replay_vs_native.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_replay_vs_native.py new file mode 100644 index 0000000000..2f3593e23e --- /dev/null +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_replay_vs_native.py @@ -0,0 +1,215 @@ +"""Isolate the `mtp_loss ~158` anomaly: is our decoupled REPLAY of the MTP head +broken, or is Megatron's NATIVE MTP forward itself wrong for Qwen3.5? + +Context: vLLM drafts this pretrained MTP head at ~0.8 acceptance (head is GOOD), yet +our training-side soft-CE reports ~158 nats (>> uniform ln(V)=12.4) -> the head looks +confidently wrong on the training side. Two hypotheses: + (a) our replay (capture mtp inputs -> re-run on detached hidden -> project) diverges + from Megatron's in-forward MTP output, or + (b) Megatron's native MTP forward is itself wrong (so even process_mtp_loss is high). + +This probe runs ONE real forward on the policy worker and, for MTP depth 0: + - captures the IN-FORWARD mtp output chunk (a forward hook) == what process_mtp_loss + projects -> in-forward logits + - runs our replay -> replay logits + - projects both through the shared output layer + - reports: max|replay - in-forward| (should be ~0 if the replay is faithful), + hard-CE of each vs ground-truth seq[t+2], and our soft-CE (teacher = rolled main). + +Interpretation: + - replay ~= in-forward AND both hard-CE high -> (b) Megatron MTP forward bug. + - replay != in-forward (replay hard-CE high, in-forward low) -> (a) replay bug. + +Run:: + uv run --isolated --extra megatron --extra dev pytest -s -vvv \ + tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_replay_vs_native.py +""" + +import pytest +import ray +import torch + +from skyrl.backends.skyrl_train.workers.megatron import ( + megatron_worker as _megatron_worker_mod, +) +from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import ( + MegatronPolicyWorkerBase, +) +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.utils import init_worker_with_type + +MODEL_NAME = "Qwen/Qwen3.5-2B" + + +class _ProbeMegatronPolicyWorker(MegatronPolicyWorkerBase): + def probe_replay_vs_native(self, token_ids=None, seq_len: int = 96) -> dict: + import torch + from megatron.core.utils import unwrap_model + + from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits + from skyrl.backends.skyrl_train.mtp.hidden_capture import ( + MTPHiddenCapture, + _resolve_mtp_host, + _unwrap_model, + ) + + gm = unwrap_model(self.actor_module[0]) + host = _resolve_mtp_host(_unwrap_model(self.actor_module[0])) + if getattr(host, "mtp", None) is None: + return {"error": "no mtp head built", "host": type(host).__name__} + + device = torch.cuda.current_device() + # Use REAL tokenized text (passed from the driver) so the model is confident -- a + # confidently-wrong MTP head shows up as high CE, which a garbage ramp cannot reveal. + if token_ids is not None: + ids = torch.tensor(token_ids, device=device, dtype=torch.long) + seq_len = ids.shape[0] + else: + ids = (torch.arange(seq_len, device=device) * 7 + 3) % 100000 + sequences = ids.unsqueeze(0) # [1, S] + attention_mask = torch.ones_like(sequences) + position_ids = torch.arange(seq_len, device=device).unsqueeze(0) + + # Capture the IN-FORWARD mtp output via a forward hook (the thing + # process_mtp_loss splits + projects), alongside our replay capture. + inforward = {} + + def _out_hook(_m, _inp, out): + inforward["out"] = out.detach() + + mtp = host.mtp + h = mtp.register_forward_hook(_out_hook) + capture = MTPHiddenCapture(self.actor_module[0], detach_trunk=True) + try: + with torch.no_grad(), capture.capture(): + outputs = self.actor_module[0](sequences, position_ids, attention_mask) + replay_hidden = capture.compute_student_hidden_states() # list per depth + finally: + h.remove() + + # Normalize a 3-D logits tensor to [B, S, V] (B=1 here). project_mtp_hidden_to_logits + # already returns [B, S, V]; the main model's output orientation is model-dependent. + def to_bsv(x): + if x.dim() != 3: + return x + if x.shape[0] == 1: + return x + if x.shape[1] == 1: + return x.transpose(0, 1).contiguous() + return x + + main_logits = to_bsv(outputs) + + # Split the in-forward output the same way compute_student_hidden_states does. + from skyrl.backends.skyrl_train.mtp.hidden_capture import _mtp_layer_offset + + num_layers = int(getattr(host.config, "mtp_num_layers", 1) or 1) + total = 1 + _mtp_layer_offset(mtp) + num_layers + inforward_hidden = list(torch.chunk(inforward["out"], total, dim=0))[-num_layers:] + + replay_logits = project_mtp_hidden_to_logits(replay_hidden, host)[0] # [B,S,V/tp] + native_logits = project_mtp_hidden_to_logits(inforward_hidden, host)[0] + + rl = to_bsv(replay_logits).float() + nl = to_bsv(native_logits).float() + ml = main_logits.float() + + # ground-truth target for depth-0 MTP: seq[t+2] + gt = torch.roll(sequences, shifts=-2, dims=1) + valid = seq_len - 2 + gt = gt[:, :valid] + + def hard_ce(logits): + lp = torch.log_softmax(logits[:, :valid].float(), dim=-1) + return -lp.gather(-1, gt.unsqueeze(-1)).squeeze(-1).mean().item() + + # soft CE: teacher = main rolled left by 1 (policy dist over seq[t+2]) + teacher = torch.roll(ml, shifts=-1, dims=1)[:, :valid] + tprob = torch.softmax(teacher, dim=-1) + r_logp = torch.log_softmax(rl[:, :valid], dim=-1) + n_logp = torch.log_softmax(nl[:, :valid], dim=-1) + soft_ce_replay = -(tprob * r_logp).sum(-1).mean().item() + soft_ce_native = -(tprob * n_logp).sum(-1).mean().item() + + replay_argmax_match = (rl[:, :valid].argmax(-1) == gt).float().mean().item() + native_argmax_match = (nl[:, :valid].argmax(-1) == gt).float().mean().item() + # Main model's own next-token (t+1) prediction -- sanity that the trunk is fine on this text. + nxt = torch.roll(sequences, -1, 1)[:, : seq_len - 1] + main_logp = torch.log_softmax(ml[:, : seq_len - 1].float(), dim=-1) + hard_ce_main_t1 = -main_logp.gather(-1, nxt.unsqueeze(-1)).squeeze(-1).mean().item() + main_argmax_match = (ml[:, : seq_len - 1].argmax(-1) == nxt).float().mean().item() + + return { + "hard_ce_main_t+1": hard_ce_main_t1, + "host": type(host).__name__, + "vocab_shard": tuple(rl.shape), + "max_abs_replay_minus_native": (rl - nl).abs().max().item(), + "hard_ce_replay": hard_ce(rl), + "hard_ce_native": hard_ce(nl), + "soft_ce_replay": soft_ce_replay, + "soft_ce_native": soft_ce_native, + "replay_argmax_match_t+2": replay_argmax_match, + "native_argmax_match_t+2": native_argmax_match, + "main_argmax_match_t+1": main_argmax_match, + "replay_logit_absmax": rl.abs().max().item(), + "native_logit_absmax": nl.abs().max().item(), + "main_logit_absmax": ml.abs().max().item(), + } + + +_ProbePolicyWorker = ray.remote(num_gpus=1)(_ProbeMegatronPolicyWorker) + + +def _cfg(): + cfg = SkyRLTrainConfig() + cfg.trainer.policy.model.path = MODEL_NAME + cfg.trainer.strategy = "megatron" + cfg.trainer.logger = "console" + cfg.trainer.placement.colocate_all = False + cfg.trainer.placement.policy_num_gpus_per_node = 1 + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.context_parallel_size = 1 + cfg.trainer.mtp.enabled = True + cfg.trainer.mtp.num_speculative_tokens = 1 + cfg.trainer.mtp.loss_type = "soft_ce" + validate_cfg(cfg) + return cfg + + +@pytest.mark.megatron +def test_mtp_replay_vs_native(ray_init_fixture): + cfg = _cfg() + + # Real coherent text so the model is confident (a confidently-wrong MTP head -> high CE). + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + passage = ( + "The mitochondria is the powerhouse of the cell. Photosynthesis converts sunlight, water, " + "and carbon dioxide into glucose and oxygen. The derivative of x squared is two x. In 1969, " + "Apollo 11 landed the first humans on the Moon. Water boils at one hundred degrees Celsius " + "at sea level. The capital of France is Paris, a city on the river Seine." + ) + token_ids = tok(passage, add_special_tokens=False)["input_ids"][:128] + + _orig = _megatron_worker_mod.PolicyWorker + _megatron_worker_mod.PolicyWorker = _ProbePolicyWorker + try: + policy = init_worker_with_type("policy", shared_pg=None, colocate_all=False, num_gpus_per_node=1, cfg=cfg) + res = ray.get(policy.async_run_ray_method("pass_through", "probe_replay_vs_native", token_ids))[0] + finally: + _megatron_worker_mod.PolicyWorker = _orig + + print("\n===== MTP replay vs native =====") + for k, v in res.items(): + print(f"{k:32s}: {v}") + print("================================") + + assert "error" not in res, res + # A faithful replay must match the in-forward MTP output. + assert res["max_abs_replay_minus_native"] < 1e-2, ( + f"replay diverges from in-forward MTP output (max_abs=" + f"{res['max_abs_replay_minus_native']:.4f}) -> replay-input bug" + ) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py new file mode 100644 index 0000000000..cad946c244 --- /dev/null +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py @@ -0,0 +1,242 @@ +"""Round-trip fidelity test for the native MTP head weights. + +Hypothesis under test: vLLM MTP speculative decoding under SkyRL gives ~0 +acceptance (vs ~0.9 when vLLM loads the HF checkpoint directly), and our +in-training ``mtp_loss`` is stuck at ~160 nats (>> uniform ln(V)=12.4). Both +point at the *weights* the MTP head sees being wrong. + +SkyRL's weight sync sends ``bridge.export_hf_weights(actor_module)`` to vLLM, +and that same loaded Megatron MTP head drives our decoupled draft replay. So if +the HF -> Megatron load -> Megatron -> HF export round-trip does NOT reproduce +the checkpoint's ``mtp.*`` tensors, that single bug explains *both* symptoms. + +This test: + 1. Builds a real Megatron *policy* worker for Qwen3.5-2B the SkyRL way + (MTP enabled), exactly as training does. + 2. Exports HF-format weights through ``bridge.export_hf_weights`` -- the + literal source of the weight sync -- and keeps the ``mtp.*`` tensors. + 3. Loads the raw ``mtp.*`` tensors straight from the HF safetensors (what a + standalone ``vllm serve`` loads, the 0.9-acceptance reference). + 4. Asserts every ``mtp.*`` tensor round-trips (shape + values close). + +Run with:: + uv run --isolated --extra megatron --extra dev pytest -s -vvv \ + tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py +""" + +import glob +import json +import os + +import pytest +import ray +import torch + +from skyrl.backends.skyrl_train.workers.megatron import ( + megatron_worker as _megatron_worker_mod, +) +from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import ( + MegatronPolicyWorkerBase, +) +from skyrl.train.config import SkyRLTrainConfig +from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.utils import init_worker_with_type + +MODEL_NAME = "Qwen/Qwen3.5-2B" + + +class _ProbeMegatronPolicyWorker(MegatronPolicyWorkerBase): + """Policy worker that exposes the exported ``mtp.*`` HF tensors (the exact + bytes the weight sync would push to vLLM), moved to CPU for the driver.""" + + def probe_export_mtp_weights(self) -> dict: + from megatron.core.utils import unwrap_model + + out = {} + all_names = [] + for name, tensor in self.bridge.export_hf_weights( + self.actor_module, + show_progress=False, + conversion_tasks=None, + ): + all_names.append(name) + # HF naming for the native MTP head is top-level ``mtp.*``; be + # liberal and also catch any nested ``.mtp.`` just in case. + if name.startswith("mtp.") or ".mtp." in name: + out[name] = tensor.detach().to(torch.float32).cpu() + del tensor + + # --- diagnostics: was the MTP head actually built on the Megatron side? --- + gm = unwrap_model(self.actor_module[0]) + # descend into VL nesting (language_model) like the capture does + host = gm + for _ in range(4): + if getattr(host, "mtp", None) is not None: + break + host = getattr(host, "language_model", None) + if host is None: + host = gm + break + built_mtp = getattr(host, "mtp", None) is not None + # collect the Megatron-side param names that mention mtp + megatron_mtp_params = [n for n, _ in gm.named_parameters() if "mtp" in n.lower()] + + return { + "mtp_tensors": out, + "total_exported": len(all_names), + "export_mtp_names": [n for n in all_names if "mtp" in n.lower()], + "sample_export_names": all_names[:5] + all_names[-5:], + "config_mtp_num_layers": getattr(getattr(host, "config", None), "mtp_num_layers", "NO_CONFIG"), + "built_mtp": built_mtp, + "host_type": type(host).__name__, + "megatron_mtp_param_names": megatron_mtp_params[:20], + "n_megatron_mtp_params": len(megatron_mtp_params), + } + + +_ProbePolicyWorker = ray.remote(num_gpus=1)(_ProbeMegatronPolicyWorker) + + +def _make_policy_cfg(model_name: str) -> SkyRLTrainConfig: + cfg = SkyRLTrainConfig() + cfg.trainer.policy.model.path = model_name + cfg.trainer.strategy = "megatron" + cfg.trainer.logger = "console" + cfg.trainer.placement.colocate_all = False + cfg.trainer.placement.policy_num_gpus_per_node = 1 + # TP=1/PP=1 so each exported tensor is whole (no sharded reassembly to muddy + # the comparison) -- this is the simplest faithful round trip. + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 + cfg.trainer.policy.megatron_config.context_parallel_size = 1 + # Enable MTP via the high-level knob, exactly as the training run does. validate_cfg -> + # _apply_mtp_config propagates this to policy.megatron_config.mtp_num_layers (setting it + # directly is clobbered back to 0/None by _apply_mtp_config when trainer.mtp.enabled is False). + cfg.trainer.mtp.enabled = True + cfg.trainer.mtp.num_speculative_tokens = 1 + cfg.trainer.mtp.loss_type = "soft_ce" + cfg.trainer.mtp.loss_weight = 0.1 + validate_cfg(cfg) + assert cfg.trainer.policy.megatron_config.mtp_num_layers == 1, ( + f"expected mtp_num_layers=1 after validate_cfg, got " + f"{cfg.trainer.policy.megatron_config.mtp_num_layers}" + ) + return cfg + + +def _load_hf_mtp_tensors(model_name: str) -> dict: + """Load every ``mtp.*`` tensor straight from the HF safetensors snapshot.""" + from huggingface_hub import snapshot_download + from safetensors import safe_open + + local = model_name + if not os.path.isdir(local): + local = snapshot_download(model_name, allow_patterns=["*.safetensors", "*.json"]) + + index = os.path.join(local, "model.safetensors.index.json") + if os.path.exists(index): + weight_map = json.load(open(index))["weight_map"] + files = sorted({os.path.join(local, f) for k, f in weight_map.items() if "mtp" in k.lower()}) + else: + files = sorted(glob.glob(os.path.join(local, "*.safetensors"))) + + out = {} + for f in files: + with safe_open(f, framework="pt", device="cpu") as handle: + for key in handle.keys(): + if key.startswith("mtp.") or ".mtp." in key: + out[key] = handle.get_tensor(key).to(torch.float32) + return out + + +def _suffix(name: str) -> str: + """Normalize to the ``mtp...`` suffix so exported and raw names align even + if one carries an extra prefix (e.g. ``model.``).""" + i = name.find("mtp.") + return name[i:] if i >= 0 else name + + +@pytest.mark.megatron +def test_mtp_head_weights_roundtrip(ray_init_fixture): + cfg = _make_policy_cfg(MODEL_NAME) + + hf = {_suffix(k): v for k, v in _load_hf_mtp_tensors(MODEL_NAME).items()} + assert hf, "no mtp.* tensors found in the HF checkpoint" + + _orig = _megatron_worker_mod.PolicyWorker + _megatron_worker_mod.PolicyWorker = _ProbePolicyWorker + try: + policy = init_worker_with_type( + "policy", + shared_pg=None, + colocate_all=False, + num_gpus_per_node=1, + cfg=cfg, + ) + probes = ray.get( + policy.async_run_ray_method("pass_through", "probe_export_mtp_weights") + ) + finally: + _megatron_worker_mod.PolicyWorker = _orig + + # --- diagnostics first --- + p0 = probes[0] + print("\n===== MTP build / export diagnostics (rank 0) =====") + print("host_type :", p0["host_type"]) + print("built_mtp (.mtp attr) :", p0["built_mtp"]) + print("config.mtp_num_layers :", p0["config_mtp_num_layers"]) + print("total exported params :", p0["total_exported"]) + print("n megatron mtp params :", p0["n_megatron_mtp_params"]) + print("megatron mtp param names:", p0["megatron_mtp_param_names"]) + print("export names w/ 'mtp' :", p0["export_mtp_names"]) + print("sample export names :", p0["sample_export_names"]) + print("==================================================") + + # TP=PP=1 -> a single rank holds the whole MTP head. + exported = {} + for rank in probes: + exported.update({_suffix(k): v for k, v in rank["mtp_tensors"].items()}) + + assert p0["built_mtp"], ( + f"MTP head was NOT built on the Megatron side (config.mtp_num_layers=" + f"{p0['config_mtp_num_layers']}, n_megatron_mtp_params={p0['n_megatron_mtp_params']}). " + "Test config did not enable MTP." + ) + assert exported, ( + "MTP head IS built on Megatron (n_megatron_mtp_params=" + f"{p0['n_megatron_mtp_params']}) but bridge.export_hf_weights exported NONE of it -> " + "the weight sync silently omits the MTP head, so vLLM never receives trained MTP weights." + ) + + print(f"\nHF mtp tensors: {len(hf)} | exported mtp tensors: {len(exported)}") + missing = sorted(set(hf) - set(exported)) + extra = sorted(set(exported) - set(hf)) + print("missing from export:", missing) + print("extra in export:", extra) + + report = [] + failures = [] + for key in sorted(hf): + h = hf[key] + if key not in exported: + failures.append(f"{key}: MISSING from export") + continue + e = exported[key] + if tuple(e.shape) != tuple(h.shape): + failures.append(f"{key}: shape {tuple(e.shape)} != HF {tuple(h.shape)}") + continue + max_abs = (e - h).abs().max().item() + denom = h.norm().item() or 1.0 + rel = (e - h).norm().item() / denom + cos = torch.nn.functional.cosine_similarity(e.flatten(), h.flatten(), dim=0).item() + report.append((key, tuple(h.shape), max_abs, rel, cos)) + # bf16 round-trip tolerance is loose; a real mismatch is off by orders + # of magnitude (rel ~ O(1), cos far from 1), not by bf16 epsilon. + if rel > 0.05 or cos < 0.99: + failures.append(f"{key}: rel={rel:.4f} cos={cos:.4f} max_abs={max_abs:.4e}") + + print("\n%-62s %-18s %10s %8s %8s" % ("name", "shape", "max_abs", "rel", "cos")) + for key, shape, max_abs, rel, cos in report: + print("%-62s %-18s %10.3e %8.4f %8.5f" % (key, str(shape), max_abs, rel, cos)) + + assert not failures, "MTP head weight round-trip mismatches:\n " + "\n ".join(failures) diff --git a/tests/backends/skyrl_train/inference_servers/test_spec_decode_drafter_reload.py b/tests/backends/skyrl_train/inference_servers/test_spec_decode_drafter_reload.py new file mode 100644 index 0000000000..041593f622 --- /dev/null +++ b/tests/backends/skyrl_train/inference_servers/test_spec_decode_drafter_reload.py @@ -0,0 +1,67 @@ +"""Unit tests for the spec-decode drafter weight-sync helper. + +Regression guard for the MTP speculative-decoding bug: vLLM's main-model weight +reload does NOT update the separate spec-decode drafter (``model_runner.drafter.model``), +so after a colocate sleep(level=2) the drafter ran on garbage and draft-acceptance +collapsed to ~0. ``_reload_spec_decode_drafter`` re-loads the drafter from the same +synced weights. These tests use lightweight stand-ins (no vLLM / GPU required). +""" + +import torch + +from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( + _reload_spec_decode_drafter, +) + + +class _RecordingDrafterModel: + def __init__(self): + self.loaded = None + + def load_weights(self, weights): + self.loaded = list(weights) + return {n for n, _ in self.loaded} + + +class _Drafter: + def __init__(self, model): + self.model = model + + +class _ModelRunner: + def __init__(self, drafter=None): + self.model = object() # the main model; irrelevant to this helper + if drafter is not None: + self.drafter = drafter + + +def _weights(): + return [("mtp.fc.weight", torch.zeros(2, 2)), ("model.embed_tokens.weight", torch.zeros(2, 2))] + + +def test_reloads_drafter_when_present(): + drafter_model = _RecordingDrafterModel() + mr = _ModelRunner(drafter=_Drafter(drafter_model)) + weights = _weights() + + assert _reload_spec_decode_drafter(mr, weights) is True + # The full weight list is forwarded; the drafter's own load_weights filters it. + assert [n for n, _ in drafter_model.loaded] == [n for n, _ in weights] + + +def test_noop_without_drafter(): + mr = _ModelRunner(drafter=None) # no spec decoding -> no .drafter attr + assert _reload_spec_decode_drafter(mr, _weights()) is False + + +def test_noop_when_drafter_has_no_model(): + mr = _ModelRunner(drafter=_Drafter(model=None)) # e.g. ngram proposer + assert _reload_spec_decode_drafter(mr, _weights()) is False + + +def test_noop_when_drafter_model_not_loadable(): + class _NoLoad: + pass + + mr = _ModelRunner(drafter=_Drafter(model=_NoLoad())) + assert _reload_spec_decode_drafter(mr, _weights()) is False diff --git a/tests/backends/skyrl_train/mtp/test_hidden_capture.py b/tests/backends/skyrl_train/mtp/test_hidden_capture.py index 59a98fa569..e82c85f11a 100644 --- a/tests/backends/skyrl_train/mtp/test_hidden_capture.py +++ b/tests/backends/skyrl_train/mtp/test_hidden_capture.py @@ -24,7 +24,9 @@ sys.modules.setdefault("megatron.core", types.ModuleType("megatron.core")) sys.modules["megatron.core.utils"] = _fake_mcore_utils -from skyrl.backends.skyrl_train.mtp.adapter import project_mtp_hidden_to_logits # noqa: E402 +from skyrl.backends.skyrl_train.mtp.adapter import ( # noqa: E402 + project_mtp_hidden_to_logits, +) from skyrl.backends.skyrl_train.mtp.hidden_capture import MTPHiddenCapture # noqa: E402 @@ -144,6 +146,100 @@ def test_project_to_logits_detach_shared_output(): assert model._out_weight.grad is None +class _FakeEmbedding(nn.Module): + """Stands in for the shared LanguageModelEmbedding the MTP block re-embeds rolled ids with.""" + + def __init__(self, vocab=5, hidden=4): + super().__init__() + self.weight = nn.Parameter(torch.randn(vocab, hidden)) + + def forward(self, input_ids=None, position_ids=None): + return self.weight[input_ids] + + +class _FakeMTPBlockWithEmbedding(_FakeMTPBlock): + """Like the real block: re-embeds rolled input ids via the ``embedding`` kwarg and mixes the + result into each depth's output, creating the second gradient path into the shared weight.""" + + def forward(self, hidden_states, input_ids=None, embedding=None, **kwargs): + emb = embedding(input_ids=input_ids).transpose(0, 1) # [b, s, h] -> [s, b, h] + chunks = [hidden_states] + [(hidden_states + emb) * self.w[i] for i in range(self.num_layers)] + return torch.cat(chunks, dim=0) + + +def _run_with_embedding(detach_shared_embedding): + model = _FakeGPT() + model.mtp = _FakeMTPBlockWithEmbedding(hidden=4, num_layers=1) + embedding = _FakeEmbedding(vocab=5, hidden=4) + capture = MTPHiddenCapture(model, detach_trunk=True, detach_shared_embedding=detach_shared_embedding) + s, b = 3, 2 + trunk = torch.randn(s, b, 4, requires_grad=True) + ids = torch.randint(0, 5, (b, s)) + with capture.capture(): + _ = model.mtp(hidden_states=trunk, input_ids=ids, embedding=embedding) + student = capture.compute_student_hidden_states() + student[0].sum().backward() + return model, embedding + + +def test_replay_detaches_shared_embedding_when_requested(): + # mtp_detach_shared_output must also sever the re-embedding path (the second route into the + # tied embedding/lm_head, next to the output projection) so only MTP-head params train. + model, embedding = _run_with_embedding(detach_shared_embedding=True) + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 + assert embedding.weight.grad is None + + +def test_replay_trains_shared_embedding_by_default(): + model, embedding = _run_with_embedding(detach_shared_embedding=False) + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 + assert embedding.weight.grad is not None and embedding.weight.grad.abs().sum() > 0 + + +def test_replay_asserts_on_positional_hidden_states(): + # The trunk detach patches kwargs only; if a future megatron passes hidden_states positionally + # the capture must fail loudly instead of silently coupling the draft loss into the trunk. + import pytest + + model = _FakeGPT() + trunk = torch.randn(3, 2, 4, requires_grad=True) + capture = MTPHiddenCapture(model, detach_trunk=True) + with capture.capture(): + _ = model.mtp(trunk, position_ids=torch.zeros(2, 3)) # positional hidden_states + with pytest.raises(AssertionError, match="keyword argument"): + capture.compute_student_hidden_states() + + +class _FakeUntiedOutputLayer(nn.Module): + def __init__(self, vocab=5, hidden=4): + super().__init__() + self.weight = nn.Parameter(torch.randn(vocab, hidden)) + + def forward(self, hidden, weight=None): + w = weight if weight is not None else self.weight + return hidden @ w.t(), None + + +class _FakeUntied(_FakeGPT): + """Untied embeddings: no shared weight; the output layer owns its own parameter.""" + + def __init__(self, hidden=4, vocab=5, mtp_num_layers=1): + super().__init__(hidden=hidden, mtp_num_layers=mtp_num_layers) + self.share_embeddings_and_output_weights = False + self.output_layer = _FakeUntiedOutputLayer(vocab=vocab, hidden=hidden) + + +def test_project_to_logits_detach_output_weight_untied_model(): + # detach_output_weight must also isolate an UNTIED model's own output-layer weight (passed + # explicitly as a detached tensor), not just the tied/shared weight. + model, trunk, student = _run(detach_trunk=True, model_cls=_FakeUntied) + logits = project_mtp_hidden_to_logits(student, model, detach_output_weight=True) + logits[0].sum().backward() + assert model.mtp.w[0].grad is not None and model.mtp.w[0].grad.abs().sum() > 0 + assert model.output_layer.weight.grad is None + assert trunk.grad is None + + class _FakeVL(nn.Module): """Stands in for a vision-language wrapper (e.g. Qwen3.5-VL ``Qwen3VLModel``): the text backbone and its MTP head are nested at ``.language_model``, and the top-level model has no ``.mtp``.""" diff --git a/tests/backends/skyrl_train/mtp/test_mtp_config.py b/tests/backends/skyrl_train/mtp/test_mtp_config.py index 49ad606e7d..32b4ca03b5 100644 --- a/tests/backends/skyrl_train/mtp/test_mtp_config.py +++ b/tests/backends/skyrl_train/mtp/test_mtp_config.py @@ -3,7 +3,14 @@ uv run --isolated --extra dev pytest tests/train/test_mtp_config.py """ -from skyrl.train.config import InferenceEngineConfig, MegatronConfig, MTPConfig, SkyRLTrainConfig +import pytest + +from skyrl.train.config import ( + InferenceEngineConfig, + MegatronConfig, + MTPConfig, + SkyRLTrainConfig, +) from skyrl.train.config.config import build_nested_dataclass from skyrl.train.utils.utils import _apply_mtp_config @@ -16,7 +23,10 @@ def test_megatron_config_mtp_defaults(): assert cfg.mtp_loss_weight == 0.1 assert cfg.mtp_loss_type == "soft_ce" assert cfg.mtp_detach_trunk is True - assert cfg.mtp_detach_shared_output is False + # Fully decoupled by default: the draft loss trains only the MTP-head parameters; the shared + # embedding/output weight (== the lm_head on tied-embedding models like Qwen3.5) is detached + # in both the output projection and the MTP block's re-embedding (slime-style). + assert cfg.mtp_detach_shared_output is True # Deprecated alias is unset by default. assert cfg.mtp_loss_scaling_factor is None @@ -67,7 +77,10 @@ def test_apply_mtp_config_enabled_propagates_to_training_and_inference(): cfg.trainer.mtp.num_speculative_tokens = 2 cfg.trainer.mtp.loss_weight = 0.25 _apply_mtp_config(cfg) - assert cfg.trainer.policy.megatron_config.mtp_num_layers == 2 + # Draft depth is inference-only: the trained head count stays None (=> the bridge infers it + # from the checkpoint), so num_speculative_tokens > 1 reuses the single head autoregressively + # in vLLM instead of force-building extra randomly-initialized Megatron heads. + assert cfg.trainer.policy.megatron_config.mtp_num_layers is None assert cfg.trainer.policy.megatron_config.mtp_loss_weight == 0.25 assert cfg.generator.inference_engine.speculative_config == { "method": "mtp", @@ -75,6 +88,28 @@ def test_apply_mtp_config_enabled_propagates_to_training_and_inference(): } +def test_apply_mtp_config_keeps_explicit_head_override(): + # A user can still pin the trained head count (e.g. force-build fresh heads on a model that + # ships without them); the draft depth stays independent. + cfg = SkyRLTrainConfig() + cfg.trainer.mtp.enabled = True + cfg.trainer.mtp.num_speculative_tokens = 3 + cfg.trainer.policy.megatron_config.mtp_num_layers = 1 + _apply_mtp_config(cfg) + assert cfg.trainer.policy.megatron_config.mtp_num_layers == 1 + assert cfg.generator.inference_engine.speculative_config["num_speculative_tokens"] == 3 + + +def test_apply_mtp_config_rejects_enabled_with_zero_heads(): + # mtp_num_layers=0 means "force-disable MTP" — contradicts trainer.mtp.enabled=true. + + cfg = SkyRLTrainConfig() + cfg.trainer.mtp.enabled = True + cfg.trainer.policy.megatron_config.mtp_num_layers = 0 + with pytest.raises(ValueError, match="mtp_num_layers=0"): + _apply_mtp_config(cfg) + + def test_apply_mtp_config_disabled_force_disables_heads(): cfg = SkyRLTrainConfig() _apply_mtp_config(cfg) diff --git a/tests/backends/skyrl_train/mtp/test_soft_ce.py b/tests/backends/skyrl_train/mtp/test_soft_ce.py index 8a7a70929e..0437d5b59a 100644 --- a/tests/backends/skyrl_train/mtp/test_soft_ce.py +++ b/tests/backends/skyrl_train/mtp/test_soft_ce.py @@ -242,6 +242,61 @@ def test_shift_mask_zeros_boundary(): assert shift_mask_for_mtp(m, 1).tolist() == [[1.0, 1.0, 0.0, 0.0]] +def test_shift_mask_left_padded_does_not_leak_pad_source(): + # Bug A regression: a left-padded row [PAD PAD t0 t1 t2]. Rolling the mask left makes the + # last pad slot (idx 1) point at the first real token, which the OLD code unmasked -> a + # de-padded zero-logit (uniform) pad position leaked into the loss. The source-side AND must + # keep only positions whose own token AND its t+shift target are real. + m = torch.tensor([[0.0, 0.0, 1.0, 1.0, 1.0]]) + # depth 0 (shift 1): valid sources are t0,t1 (targets t1,t2 real); t2 has no real target. + assert shift_mask_for_mtp(m, 0).tolist() == [[0.0, 0.0, 1.0, 1.0, 0.0]] + # depth 1 (shift 2): only t0 has a real target (t2); pad idx0/idx1 must stay 0. + assert shift_mask_for_mtp(m, 1).tolist() == [[0.0, 0.0, 1.0, 0.0, 0.0]] + + +def test_shift_mask_right_padded_is_unaffected(): + # Right padding never leaks (rolled mask is already a subset), so the fix is a no-op there. + m = torch.tensor([[1.0, 1.0, 1.0, 0.0, 0.0]]) + assert shift_mask_for_mtp(m, 0).tolist() == [[1.0, 1.0, 0.0, 0.0, 0.0]] + + +def test_left_pad_zero_logits_do_not_inflate_loss(): + # End-to-end (loss-level) reproduction of Bug A: emulate the de-pad pipeline, which ZERO-fills + # pad positions (postprocess_packed_seqs / recover_left_padding both use torch.zeros). With a + # perfectly-aligned student at the real positions, the draft soft-CE must equal the teacher's + # entropy over the real supervised positions -- the leaked uniform pad position must NOT inflate + # it. Also asserts the leak (had it survived) is bounded by log(V), per the two-bug analysis. + torch.manual_seed(0) + V = 64 + pad = 2 # left padding + real = 6 + seq = pad + real + # Real-position teacher logits: moderately peaked so entropy << log(V). + main_logits = torch.zeros(1, seq, V) + main_logits[:, pad:, :] = torch.randn(1, real, V) * 3.0 # pad positions stay zero (de-pad fill) + mask = torch.zeros(1, seq) + mask[:, pad:] = 1.0 + + # Perfectly-aligned student: student[t] == teacher target for depth 0 == main_logits[t+1]. + # Build it from the rolled teacher so soft-CE at real+aligned positions == teacher entropy. + teacher = build_teacher_logits(main_logits, 0) # roll(-1); pad positions stay zero + student = teacher.clone() # aligned where it matters; pad positions zero (uniform), like de-pad + + layer_mask = shift_mask_for_mtp(mask, 0) + loss = draft_soft_ce(student, teacher, layer_mask) + + # Oracle: entropy of the teacher over exactly the source-AND-target-valid positions. + valid = (mask[:, :].bool()) & (torch.roll(mask, -1, 1).bool()) + valid[:, -1:] = False + tprob = F.softmax(teacher.float(), dim=-1) + ent = -(tprob * torch.log_softmax(teacher.float(), -1)).sum(-1) + oracle = (ent * valid).sum() / valid.sum() + assert torch.allclose(loss, oracle, atol=1e-5), (loss.item(), oracle.item()) + # And the result is the true (low) entropy, nowhere near log(V). + assert loss.item() < ent[valid].max().item() + 1e-4 + assert loss.item() < torch.log(torch.tensor(float(V))).item() + + def test_combine_is_policy_plus_weighted_draft(): torch.manual_seed(3) policy_loss = torch.tensor(2.0, requires_grad=True) @@ -250,9 +305,7 @@ def test_combine_is_policy_plus_weighted_draft(): mask = torch.ones(2, 5) cfg = DraftLossConfig(loss_weight=0.5, loss_type="soft_ce") - combined, metrics = combine_policy_and_draft_loss( - policy_loss, [student], main_logits, mask, cfg - ) + combined, metrics = combine_policy_and_draft_loss(policy_loss, [student], main_logits, mask, cfg) # Recompute the draft term independently and check the combination identity. teacher = build_teacher_logits(main_logits, 0) draft = draft_soft_ce(student, teacher, shift_mask_for_mtp(mask, 0)) @@ -269,9 +322,64 @@ def test_combine_hard_ce_uses_rolled_labels(): base_labels = torch.randint(0, 7, (1, 5)) cfg = DraftLossConfig(loss_weight=1.0, loss_type="hard_ce") - combined, _ = combine_policy_and_draft_loss( - policy_loss, [student], main_logits, mask, cfg, hard_labels=base_labels - ) + combined, _ = combine_policy_and_draft_loss(policy_loss, [student], main_logits, mask, cfg, hard_labels=base_labels) layer_labels = torch.roll(base_labels, shifts=-1, dims=1) - draft = draft_hard_ce(student, layer_labels, shift_mask_for_mtp(mask, 0)) + # Hard CE masks with one extra shift (label at t+k+2 must be a real token), unlike soft CE. + draft = draft_hard_ce(student, layer_labels, shift_mask_for_mtp(mask, 1)) assert torch.allclose(combined, policy_loss + draft, atol=1e-6) + + +def test_combine_hard_ce_excludes_boundary_with_invalid_label(): + # Regression: hard CE at depth k supervises position t against the TOKEN seq[t+k+2] — one + # position further than the soft-CE teacher (a distribution read at t+k+1). With the soft-CE + # mask shift, the last in-range position (t = last_real - (k+1)) was trained against the + # zeroed boundary label (token id 0). Build a student that is perfectly confident in the true + # label at every genuinely-valid position but puts ~zero mass on token 0 at the boundary + # position: the loss must be ~0 (boundary excluded), not inflated by -log q(0). + S, V = 6, 7 + seq = torch.arange(1, S + 1) % V # tokens 1..6 — never 0, so label 0 only comes from the boundary + hard_labels = torch.roll(seq, -1).unsqueeze(0).clone() # labels[t] = seq[t+1] + hard_labels[:, -1] = 0 + mask = torch.ones(1, S) + + # depth 0: true label at t is seq[t+2]; valid sources are t = 0..S-3. + student = torch.full((1, S, V), -30.0) + for t in range(S - 2): + student[0, t, seq[t + 2]] = 30.0 + # boundary position t = S-2: confident in a NON-zero token, so a leaked label-0 target + # would contribute ~60 nats. + student[0, S - 2, 1] = 30.0 + + cfg = DraftLossConfig(loss_weight=1.0, loss_type="hard_ce") + combined, metrics = combine_policy_and_draft_loss( + torch.tensor(0.0), [student], torch.zeros(1, S, V), mask, cfg, hard_labels=hard_labels + ) + assert metrics["draft_loss"] < 1e-3, metrics["draft_loss"] + + +def test_unpadded_vocab_shard_width(): + from skyrl.backends.skyrl_train.mtp.soft_ce import unpadded_vocab_shard_width + + # Unknown true vocab -> no-op. + assert unpadded_vocab_shard_width(None, 128, 0) == 128 + # No padding: every rank keeps its full shard. + assert unpadded_vocab_shard_width(256, 128, 0) == 128 + assert unpadded_vocab_shard_width(256, 128, 1) == 128 + # 250 padded to 256 over TP=2: rank 0 full, rank 1 loses the 6-column tail. + assert unpadded_vocab_shard_width(250, 128, 0) == 128 + assert unpadded_vocab_shard_width(250, 128, 1) == 122 + # Pathological: a rank whose entire shard is padding. + assert unpadded_vocab_shard_width(100, 128, 1) == 0 + + +def test_draft_soft_ce_topk_grad_dtype_matches_input(): + # The top-k backward must build its full-vocab grad buffer directly in the input dtype + # (bf16), not fp32-then-cast — value-identical, half the transient at large vocab. + from skyrl.backends.skyrl_train.mtp.soft_ce import draft_soft_ce_topk + + torch.manual_seed(0) + student = torch.randn(1, 4, 9, dtype=torch.bfloat16, requires_grad=True) + teacher = torch.randn(1, 4, 9, dtype=torch.bfloat16) + draft_soft_ce_topk(student, teacher, torch.ones(1, 4), k=3).backward() + assert student.grad.dtype == torch.bfloat16 + assert int((student.grad != 0).sum(-1).max()) <= 3 diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_client.py b/tests/backends/skyrl_train/mtp/test_spec_decode_client.py index 85269cd9d9..1c707c136e 100644 --- a/tests/backends/skyrl_train/mtp/test_spec_decode_client.py +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_client.py @@ -82,8 +82,6 @@ async def test_get_spec_decode_metrics_returns_none_when_no_engine_reports(): @pytest.mark.asyncio async def test_get_spec_decode_metrics_ignores_non_reporting_engines(): # A mix of reporting and non-reporting engines still aggregates the reporters only. - client = _make_spec_client( - [None, {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5}] - ) + client = _make_spec_client([None, {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5}]) totals = await client.get_spec_decode_metrics() assert totals == {"num_drafts": 2, "num_draft_tokens": 8, "num_accepted_tokens": 5} diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py index 4d2c15408d..012602558a 100644 --- a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py @@ -5,23 +5,37 @@ import types +import pytest + from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( acceptance_rate_metrics, make_spec_decode_stat_logger_class, + merge_spec_decode_counters, sum_spec_decode_loggers, ) -def _sched(num_draft, num_accepted, num_drafts): +def _logger_cls(): + # The stat-logger class factory lazily imports vLLM (it subclasses StatLoggerBase); skip the + # logger-class tests in vLLM-less environments instead of failing. The pure aggregation / + # rate-math tests below run everywhere. + pytest.importorskip("vllm") + return make_spec_decode_stat_logger_class() + + +def _sched(num_draft, num_accepted, num_drafts, per_pos=None): return types.SimpleNamespace( spec_decoding_stats=types.SimpleNamespace( - num_drafts=num_drafts, num_draft_tokens=num_draft, num_accepted_tokens=num_accepted + num_drafts=num_drafts, + num_draft_tokens=num_draft, + num_accepted_tokens=num_accepted, + num_accepted_tokens_per_pos=per_pos if per_pos is not None else [], ) ) def test_logger_accumulates_across_iterations(): - cls = make_spec_decode_stat_logger_class() + cls = _logger_cls() lg = cls(vllm_config=None, engine_index=0) lg.record(scheduler_stats=_sched(10, 7, 5)) lg.record(scheduler_stats=_sched(8, 6, 4)) @@ -29,39 +43,123 @@ def test_logger_accumulates_across_iterations(): assert (lg.num_draft_tokens, lg.num_accepted_tokens, lg.num_drafts) == (18, 13, 9) +def test_logger_accumulates_per_position_and_grows(): + # The per-position list adapts to whatever depth vLLM reports (here it grows 2 -> 3, e.g. the + # configured num_speculative_tokens), with missing tail positions treated as 0. + cls = _logger_cls() + lg = cls(vllm_config=None, engine_index=0) + lg.record(scheduler_stats=_sched(10, 7, 5, per_pos=[5, 2])) + lg.record(scheduler_stats=_sched(12, 9, 4, per_pos=[4, 3, 2])) + lg.record(scheduler_stats=types.SimpleNamespace(spec_decoding_stats=None)) + assert lg.num_accepted_tokens_per_pos == [9, 5, 2] + + +def test_logger_tolerates_stats_without_per_position(): + # Older vLLM stats objects without num_accepted_tokens_per_pos must not break accounting. + cls = _logger_cls() + lg = cls(vllm_config=None, engine_index=0) + stats = types.SimpleNamespace( + spec_decoding_stats=types.SimpleNamespace(num_drafts=3, num_draft_tokens=6, num_accepted_tokens=4) + ) + lg.record(scheduler_stats=stats) + assert (lg.num_draft_tokens, lg.num_accepted_tokens, lg.num_drafts) == (6, 4, 3) + assert lg.num_accepted_tokens_per_pos == [] + + def test_sum_spec_decode_loggers(): - cls = make_spec_decode_stat_logger_class() - a = cls(None, 0) - a.num_draft_tokens, a.num_accepted_tokens, a.num_drafts = 10, 6, 3 - b = cls(None, 1) - b.num_draft_tokens, b.num_accepted_tokens, b.num_drafts = 4, 2, 1 + # sum_spec_decode_loggers only reads counter attributes, so plain namespaces stand in for + # logger instances and the test runs without vLLM. + a = types.SimpleNamespace( + num_draft_tokens=10, num_accepted_tokens=6, num_drafts=3, num_accepted_tokens_per_pos=[4, 2] + ) + b = types.SimpleNamespace( + num_draft_tokens=4, num_accepted_tokens=2, num_drafts=1, num_accepted_tokens_per_pos=[1, 1, 0] + ) assert sum_spec_decode_loggers([a, b]) == { "num_drafts": 4, "num_draft_tokens": 14, "num_accepted_tokens": 8, + "num_accepted_tokens_per_pos": [5, 3, 0], } assert sum_spec_decode_loggers([]) is None +def test_merge_spec_decode_counters_across_engines(): + totals = {} + merge_spec_decode_counters( + totals, + {"num_drafts": 3, "num_draft_tokens": 9, "num_accepted_tokens": 6, "num_accepted_tokens_per_pos": [3, 2, 1]}, + ) + merge_spec_decode_counters( + totals, + {"num_drafts": 2, "num_draft_tokens": 4, "num_accepted_tokens": 3, "num_accepted_tokens_per_pos": [2, 1]}, + ) + assert totals == { + "num_drafts": 5, + "num_draft_tokens": 13, + "num_accepted_tokens": 9, + "num_accepted_tokens_per_pos": [5, 3, 1], + } + + def test_acceptance_rate_per_step_delta(): # cumulative counters grow each step; the metric is the per-step delta ratio. prev = None - m1, prev = acceptance_rate_metrics( - {"num_drafts": 5, "num_draft_tokens": 10, "num_accepted_tokens": 6}, prev - ) + m1, prev = acceptance_rate_metrics({"num_drafts": 5, "num_draft_tokens": 10, "num_accepted_tokens": 6}, prev) assert m1["vllm/draft_num_draft_tokens"] == 10 assert m1["vllm/draft_num_accepted_tokens"] == 6 assert abs(m1["vllm/draft_acceptance_rate"] - 0.6) < 1e-9 - m2, prev = acceptance_rate_metrics( - {"num_drafts": 11, "num_draft_tokens": 22, "num_accepted_tokens": 15}, prev - ) + m2, prev = acceptance_rate_metrics({"num_drafts": 11, "num_draft_tokens": 22, "num_accepted_tokens": 15}, prev) # deltas: drafted 22-10=12, accepted 15-6=9 -> 0.75 assert m2["vllm/draft_num_draft_tokens"] == 12 assert m2["vllm/draft_num_accepted_tokens"] == 9 assert abs(m2["vllm/draft_acceptance_rate"] - 0.75) < 1e-9 +def test_acceptance_rate_per_position(): + # Per-position rate = drafts this step whose k-th speculated token was accepted / drafts this + # step (1-based keys). One key per configured draft position; non-increasing in k. + prev = None + m1, prev = acceptance_rate_metrics( + { + "num_drafts": 10, + "num_draft_tokens": 30, + "num_accepted_tokens": 17, + "num_accepted_tokens_per_pos": [8, 6, 3], + }, + prev, + ) + assert abs(m1["vllm/draft_acceptance_rate_pos_1"] - 0.8) < 1e-9 + assert abs(m1["vllm/draft_acceptance_rate_pos_2"] - 0.6) < 1e-9 + assert abs(m1["vllm/draft_acceptance_rate_pos_3"] - 0.3) < 1e-9 + + # Second step: only the deltas count (drafts 10 -> 14, pos counts [8,6,3] -> [11,8,4]). + m2, prev = acceptance_rate_metrics( + { + "num_drafts": 14, + "num_draft_tokens": 42, + "num_accepted_tokens": 24, + "num_accepted_tokens_per_pos": [11, 8, 4], + }, + prev, + ) + assert abs(m2["vllm/draft_acceptance_rate_pos_1"] - 3 / 4) < 1e-9 + assert abs(m2["vllm/draft_acceptance_rate_pos_2"] - 2 / 4) < 1e-9 + assert abs(m2["vllm/draft_acceptance_rate_pos_3"] - 1 / 4) < 1e-9 + + +def test_acceptance_rate_per_position_depth_one_single_key(): + # Depth 1 (the current default) emits exactly one per-position key. + metrics, _ = acceptance_rate_metrics( + {"num_drafts": 4, "num_draft_tokens": 4, "num_accepted_tokens": 3, "num_accepted_tokens_per_pos": [3]}, + None, + ) + pos_keys = [k for k in metrics if "_pos_" in k] + assert pos_keys == ["vllm/draft_acceptance_rate_pos_1"] + assert abs(metrics["vllm/draft_acceptance_rate_pos_1"] - 0.75) < 1e-9 + + def test_acceptance_rate_no_spec_decode_returns_empty(): metrics, prev = acceptance_rate_metrics(None, None) assert metrics == {} @@ -70,7 +168,5 @@ def test_acceptance_rate_no_spec_decode_returns_empty(): def test_acceptance_rate_no_drafts_omits_rate(): # zero drafted tokens -> report counts but no (undefined) rate. - metrics, _ = acceptance_rate_metrics( - {"num_drafts": 0, "num_draft_tokens": 0, "num_accepted_tokens": 0}, None - ) + metrics, _ = acceptance_rate_metrics({"num_drafts": 0, "num_draft_tokens": 0, "num_accepted_tokens": 0}, None) assert metrics == {"vllm/draft_num_draft_tokens": 0, "vllm/draft_num_accepted_tokens": 0} diff --git a/tests/backends/skyrl_train/workers/test_worker_utils.py b/tests/backends/skyrl_train/workers/test_worker_utils.py index 2c01499faa..5b082367ae 100644 --- a/tests/backends/skyrl_train/workers/test_worker_utils.py +++ b/tests/backends/skyrl_train/workers/test_worker_utils.py @@ -59,6 +59,19 @@ def test_reduce_metrics_mixed(self): assert result["policy_loss"] == 4.0 # sum assert result["entropy"] == 2.0 # mean + def test_reduce_metrics_mtp_loss_is_averaged_not_summed(self): + """mtp_loss / draft_loss are per-token MEANS, not pre-scaled sums. They must be averaged + across microbatches even when sum_loss_metrics=True, or a true ~0.5 reads as ~N*0.5.""" + metrics = { + "policy_loss": [1.0, 3.0], # pre-scaled -> sum + "mtp_loss": [0.5, 0.5, 0.5, 0.5], # mean -> stays 0.5, NOT 2.0 + "draft_loss": [0.6, 0.4], # mean -> 0.5 + } + result = reduce_metrics(metrics, sum_loss_metrics=True) + assert result["policy_loss"] == 4.0 # still summed + assert result["mtp_loss"] == 0.5 # averaged, not 2.0 + assert result["draft_loss"] == 0.5 + def test_reduce_metrics_single_value(self): """Test reduction with single value lists.""" metrics = { From 2e824f0a5ee0c8627144f6798843622fe1a90c98 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 15 Jun 2026 17:35:03 +0000 Subject: [PATCH 13/28] fix trainer slow bug --- .../megatron/run_megatron_dapo_qwen3.5_9b.sh | 8 +- ...run_megatron_dapo_qwen3.5_9b_specdecode.sh | 199 ++++++++++++++++ ...tron_dapo_qwen3.5_9b_specdecode_notrain.sh | 213 ++++++++++++++++++ skyrl/backends/skyrl_train/mtp/adapter.py | 31 +++ .../megatron/megatron_model_wrapper.py | 60 ++++- skyrl/train/utils/utils.py | 1 + .../backends/skyrl_train/mtp/test_adapter.py | 78 +++++++ 7 files changed, 577 insertions(+), 13 deletions(-) create mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh create mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh create mode 100644 tests/backends/skyrl_train/mtp/test_adapter.py diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh index 446da8a885..042ef75985 100644 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh @@ -21,8 +21,8 @@ NUM_GPUS_PER_NODE=8 # 9B is ~4.5x the 2B: a single full-weight copy is ~18GB in bf16. Use inference # TP=2 (4 engines) so each engine's weights are halved (~9GB/GPU) and there is # more headroom for KV cache during generation. TP=2 comm stays on NVLink. -NUM_INFERENCE_ENGINES=4 -INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=2 +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 LOGGER="wandb" # change to "console" to print to stdout CLIP_RATIO_LOW=0.2 @@ -45,7 +45,7 @@ MAX_PROMPT_LENGTH=$((1024 * 2)) MAX_RESPONSE_LENGTH=$((1024 * 8)) # repro run parameters -TRAIN_BATCH_SIZE=128 +TRAIN_BATCH_SIZE=32 MINI_BATCH_SIZE=32 N_SAMPLES_PER_PROMPT=8 EVAL_N_SAMPLES_PER_PROMPT=16 @@ -147,7 +147,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ generator.inference_engine.gpu_memory_utilization=0.5 \ trainer.logger="$LOGGER" \ trainer.project_name="qwen3_5_dapo" \ - trainer.run_name="dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.run_name="nosd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.export_path="/mnt/local_storage/exports/dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.hf_save_interval=300 \ trainer.resume_mode=latest \ diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh new file mode 100644 index 0000000000..cb19c3f6ea --- /dev/null +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh @@ -0,0 +1,199 @@ +set -x + +# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO with Megatron, +# WITH Multi-Token Prediction (MTP) speculative decoding for faster rollout. +# +# This is the spec-decode counterpart of run_megatron_dapo_qwen3.5_9b.sh. Every +# knob below is IDENTICAL to the no-spec script (same batch sizes, LR, parallelism, +# sampling) so a reward-curve / throughput comparison is apples-to-apples. The ONLY +# difference is the `trainer.mtp.*` block at the bottom. +# +# What MTP on does (single high-level `trainer.mtp` knob, see skyrl/train/config/config.py): +# - Training side: builds + trains Qwen3.5-9B's native MTP head (the model ships +# `mtp_num_hidden_layers: 1`) with a decoupled draft loss (top-k soft-CE distillation against +# the policy's own detached next-token distribution). The draft gradient never pulls on +# the policy backbone. +# - Inference side: enables vLLM MTP speculative decoding +# (`speculative_config={"method": "mtp", "num_speculative_tokens": 3}`). With a single trained +# head, depth>1 reuses that head autoregressively at draft time. vLLM loads the MTP head from the +# same policy checkpoint, and SkyRL's weight sync keeps the draft head in sync with the trained +# policy each step. +# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. +# +# Runs on 1 node of 8xH100s (80GB each). +# +# NOTE: verify the exact HF repo id for the 9B model before running +# (e.g. `hf download Qwen/Qwen3.5-9B` / check https://huggingface.co/Qwen). +# +# Prepare data onto the fast local disk first: +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# Then launch: +# bash examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh + +MODEL_NAME="Qwen/Qwen3.5-9B" +# Use the fast, non-persistent local disk for data (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +# 9B is ~4.5x the 2B: a single full-weight copy is ~18GB in bf16. Use inference +# TP=2 (4 engines) so each engine's weights are halved (~9GB/GPU) and there is +# more headroom for KV cache during generation. TP=2 comm stays on NVLink. +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 +LOGGER="wandb" # change to "console" to print to stdout + +CLIP_RATIO_LOW=0.2 +CLIP_RATIO_HIGH=0.28 +# use token mean loss reduction +LOSS_REDUCTION="token_mean" +# applies overlong filtering (but not soft overlong punishment) +APPLY_OVERLONG_FILTERING=true +# apply soft overlong punishment with custom trainer impl in main_dapo.py +OVERLONG_BUFFER_LEN=$((1024 * 4)) +OVERLONG_BUFFER_PENALTY_FACTOR=1.0 + +# other DAPO parameters +USE_KL_LOSS=false +TEMPERATURE=1.0 +TOP_P=1.0 +EVAL_TOP_P=0.7 +CLIP_RATIO_C=10.0 +MAX_PROMPT_LENGTH=$((1024 * 2)) +MAX_RESPONSE_LENGTH=$((1024 * 8)) + +# repro run parameters +TRAIN_BATCH_SIZE=32 +MINI_BATCH_SIZE=32 +N_SAMPLES_PER_PROMPT=8 +EVAL_N_SAMPLES_PER_PROMPT=16 +ENFORCE_EAGER=true # cuda graphs can cause some instability +LR=1e-6 + +# megatron config -- Qwen3.5-9B is a dense model, so no expert parallelism. +# TP=4 (up from 2 on the 2B): 9B params + Adam states + 8K-token activations +# need more sharding to fit at micro batch 1. TP>1 auto-enables sequence +# parallelism, sharding activations/vocab-logits across the TP group. +# TP=4, PP=1, CP=1 => DP=2. TP stays within the single-node NVLink domain. +MEGATRON_TP=4 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + +# optimizer offload -- OFF by default (with colocate_all the inference engine is +# offloaded during training, so the full 80GB is available for the optimizer). +# Flip to true if you hit OOM in the optimizer step / grad sync. +OPTIMIZER_OFFLOAD=false +OPTIMIZER_OFFLOAD_FRACTION=1.0 + +# TIS parameters +TIS_IMP_RATIO_CAP=2.0 +TIS_TYPE=token + + +# Multi-Token Prediction (MTP) speculative decoding. +# Qwen3.5-9B ships 1 native MTP head (`mtp_num_hidden_layers: 1`); training always trains the +# checkpoint's heads. NUM_SPECULATIVE_TOKENS is the vLLM *draft depth* only — values > 1 reuse the +# single head autoregressively at draft time (more speedup, but per-position acceptance decays with +# depth since the head never trains on its own outputs). Here k=3. +MTP_ENABLED=true +MTP_NUM_SPECULATIVE_TOKENS=3 +MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) +MTP_LOSS_WEIGHT=0.5 +# NOTE: trainer.policy.megatron_config.mtp_detach_shared_output now defaults to true: the draft loss +# trains ONLY the MTP-head params; the tied embedding/lm_head is detached (output projection AND the +# MTP block's re-embedding), so the draft gradient no longer nudges the policy's own logits. Set it +# to false for the old NeMo-RL behaviour (shared head also trained by the draft loss). +# Top-k draft loss: distill only the teacher's top-k tokens instead of the full 248K vocab, keeping +# draft-loss memory at O(seq*k) vs O(seq*vocab). Only applies to mtp_loss_type="soft_ce". Here k=256. +MTP_LOSS_TOPK=256 + + +# Qwen3.5 flags +REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 +ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 +DISTRIBUTED_EXECUTOR_BACKEND="mp" +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="dual_clip" \ + trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ + trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + generator.sampling_params.temperature=$TEMPERATURE \ + generator.sampling_params.top_p=$TOP_P \ + generator.eval_sampling_params.top_p=$EVAL_TOP_P \ + generator.eval_sampling_params.temperature=$TEMPERATURE \ + generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=10 \ + trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ + trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=-1 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=$LR \ + trainer.policy.optimizer_config.num_warmup_steps=5 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.mtp.enabled=$MTP_ENABLED \ + trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ + trainer.mtp.loss_type=$MTP_LOSS_TYPE \ + trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ + trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ + trainer.logger="$LOGGER" \ + trainer.project_name="qwen3_5_dapo" \ + trainer.run_name="sd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.export_path="/mnt/local_storage/exports/sd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.hf_save_interval=300 \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_path="/mnt/local_storage/ckpts/sd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + $@ diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh new file mode 100644 index 0000000000..b5a6731ab4 --- /dev/null +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh @@ -0,0 +1,213 @@ +set -x + +# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO with Megatron, +# WITH vLLM MTP speculative decoding for faster rollout, but the MTP head is NOT trained. +# +# This is the "frozen-head" counterpart of run_megatron_dapo_qwen3.5_9b_specdecode.sh. +# Every knob is IDENTICAL except the MTP block at the bottom: the head is BUILT and SYNCED +# like the training variant, but its draft loss weight is 0 so it is never actually trained. +# +# Why not just turn the head off on the training side? Under colocate_all, vLLM is slept +# with level=2 between steps, which DISCARDS its weights -- including the separate spec-decode +# drafter (the MTP head). On wake, SkyRL restores the drafter ONLY from the Megatron +# weight-sync list (inference_servers/spec_decode_utils.py::_reload_spec_decode_drafter). If +# Megatron doesn't build the head, the sync list has no `mtp.*` weights, the drafter stays +# garbage, and draft acceptance is ~0 (many drafted tokens, almost all rejected). So the head +# MUST be built + exported every sync even when we don't want to train it. +# +# What this configuration does: +# - Training side: `trainer.mtp.enabled=true` builds Qwen3.5-9B's native MTP head (the model +# ships `mtp_num_hidden_layers: 1`) and exports it through weight sync each step. But +# `trainer.mtp.loss_weight=0.0` makes the decoupled draft loss contribute ZERO gradient, so +# the head stays at its pretrained values (only ~1e-7/step weight-decay drift). The RL policy +# backbone trains normally. +# - Inference side: vLLM MTP speculative decoding is enabled with `num_speculative_tokens` +# draft tokens; the drafter is re-synced with the (frozen) pretrained head every step. +# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. +# +# CAVEAT: the frozen head is matched to the ORIGINAL base model. As RL updates the policy +# backbone (and syncs those updates to vLLM), the never-updated head drifts out of alignment, +# so `vllm/draft_acceptance_rate` will DECAY over training -- best early, eroding as the policy +# moves. Training the head (the _specdecode.sh variant) closes that gap. +# +# Runs on 1 node of 8xH100s (80GB each). +# +# NOTE: verify the exact HF repo id for the 9B model before running +# (e.g. `hf download Qwen/Qwen3.5-9B` / check https://huggingface.co/Qwen). +# +# Prepare data onto the fast local disk first: +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# Then launch: +# bash examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh + +MODEL_NAME="Qwen/Qwen3.5-9B" +# Use the fast, non-persistent local disk for data (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +# 9B is ~4.5x the 2B: a single full-weight copy is ~18GB in bf16. Use inference +# TP=2 (4 engines) so each engine's weights are halved (~9GB/GPU) and there is +# more headroom for KV cache during generation. TP=2 comm stays on NVLink. +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 +LOGGER="wandb" # change to "console" to print to stdout + +CLIP_RATIO_LOW=0.2 +CLIP_RATIO_HIGH=0.28 +# use token mean loss reduction +LOSS_REDUCTION="token_mean" +# applies overlong filtering (but not soft overlong punishment) +APPLY_OVERLONG_FILTERING=true +# apply soft overlong punishment with custom trainer impl in main_dapo.py +OVERLONG_BUFFER_LEN=$((1024 * 4)) +OVERLONG_BUFFER_PENALTY_FACTOR=1.0 + +# other DAPO parameters +USE_KL_LOSS=false +TEMPERATURE=1.0 +TOP_P=1.0 +EVAL_TOP_P=0.7 +CLIP_RATIO_C=10.0 +MAX_PROMPT_LENGTH=$((1024 * 2)) +MAX_RESPONSE_LENGTH=$((1024 * 8)) + +# repro run parameters +TRAIN_BATCH_SIZE=32 +MINI_BATCH_SIZE=32 +N_SAMPLES_PER_PROMPT=8 +EVAL_N_SAMPLES_PER_PROMPT=16 +ENFORCE_EAGER=true # cuda graphs can cause some instability +LR=1e-6 + +# megatron config -- Qwen3.5-9B is a dense model, so no expert parallelism. +# TP=4 (up from 2 on the 2B): 9B params + Adam states + 8K-token activations +# need more sharding to fit at micro batch 1. TP>1 auto-enables sequence +# parallelism, sharding activations/vocab-logits across the TP group. +# TP=4, PP=1, CP=1 => DP=2. TP stays within the single-node NVLink domain. +MEGATRON_TP=4 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + +# optimizer offload -- OFF by default (with colocate_all the inference engine is +# offloaded during training, so the full 80GB is available for the optimizer). +# Flip to true if you hit OOM in the optimizer step / grad sync. +OPTIMIZER_OFFLOAD=false +OPTIMIZER_OFFLOAD_FRACTION=1.0 + +# TIS parameters +TIS_IMP_RATIO_CAP=2.0 +TIS_TYPE=token + + +# vLLM MTP speculative decoding WITHOUT training the head (frozen pretrained head baseline). +# +# IMPORTANT: under colocate_all, vLLM is slept with level=2 between steps, which DISCARDS its +# weights — including the separate spec-decode drafter (the MTP head). On wake, SkyRL restores the +# drafter ONLY from the Megatron weight-sync list (see inference_servers/spec_decode_utils.py: +# _reload_spec_decode_drafter). So the head MUST be built + exported on the Megatron side every sync, +# or the drafter stays garbage and draft acceptance is ~0 (lots of drafted tokens, almost all +# rejected). vLLM's init-time checkpoint load is NOT enough — sleep(level=2) throws it away on step 1. +# +# Therefore we keep trainer.mtp.enabled=true (builds the native head, exports it each sync, sets +# vLLM speculative_config) but set loss_weight=0.0 so the draft loss contributes ZERO gradient: the +# head stays at its pretrained values (only ~1e-7/step weight-decay drift, negligible) and is +# re-synced to vLLM every step. Acceptance then reflects the genuine "stock pretrained head, drifting +# RL backbone" baseline. NOTE: the (top-k soft-CE) draft loss is still computed each step and thrown +# away — small wasted compute. Ask for the proper freeze path (skip the loss + freeze the params) if +# that matters. +MTP_NUM_SPECULATIVE_TOKENS=3 +MTP_LOSS_TYPE="soft_ce" # kept cheap via top-k; weight is 0 so the value is discarded +MTP_LOSS_WEIGHT=0.0 # 0 => head is built + synced but never trained (frozen pretrained head) +MTP_LOSS_TOPK=256 + + +# Qwen3.5 flags +REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 +ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 +DISTRIBUTED_EXECUTOR_BACKEND="mp" +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="dual_clip" \ + trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ + trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ + trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ + generator.sampling_params.temperature=$TEMPERATURE \ + generator.sampling_params.top_p=$TOP_P \ + generator.eval_sampling_params.top_p=$EVAL_TOP_P \ + generator.eval_sampling_params.temperature=$TEMPERATURE \ + generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ + trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \ + trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=10 \ + trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ + trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ + trainer.eval_batch_size=1024 \ + trainer.eval_before_train=false \ + trainer.eval_interval=5 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=$TRAIN_BATCH_SIZE \ + trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ + trainer.micro_forward_batch_size_per_gpu=1 \ + trainer.micro_train_batch_size_per_gpu=1 \ + trainer.ckpt_interval=50 \ + trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ + generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ + trainer.policy.optimizer_config.lr=$LR \ + trainer.policy.optimizer_config.num_warmup_steps=5 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.async_engine=true \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ + generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.mtp.enabled=true \ + trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ + trainer.mtp.loss_type=$MTP_LOSS_TYPE \ + trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ + trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ + trainer.logger="$LOGGER" \ + trainer.project_name="qwen3_5_dapo" \ + trainer.run_name="sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.export_path="/mnt/local_storage/exports/sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + trainer.hf_save_interval=300 \ + trainer.resume_mode=latest \ + trainer.max_ckpts_to_keep=3 \ + trainer.ckpt_path="/mnt/local_storage/ckpts/sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ + $@ diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py index 085197376a..db8750230a 100644 --- a/skyrl/backends/skyrl_train/mtp/adapter.py +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -20,6 +20,36 @@ from typing import List, Protocol, runtime_checkable +import torch + + +class _CanonicalGradStrides(torch.autograd.Function): + """Identity forward; backward hands upstream a stride-canonical gradient. + + The draft projection runs through megatron's ``LinearWithFrozenWeight`` (the + detached-weight path), whose hand-written dgrad is ``grad_output.matmul(weight)``. + ``torch.matmul`` only folds that 3D x 2D product into a flat ``mm`` when the grad + passes ``should_fold``'s stride check, which does NOT skip size-1 dims. With + micro-batch 1, the grad reaching it is a transpose-backward view whose size-1 batch + dim carries a stale stride (the adapter's ``.contiguous()`` no-ops on such views), + so matmul falls back to a broadcast ``bmm`` (batch=seq, M=1) that runs ~100x slower + than the equivalent ``mm`` — measured 1.27s vs 10ms per microbatch at + [8344, 1, 62080] x [62080, 4096] on H100. Re-viewing the (dense) grad rewrites the + strides to canonical form at zero copy and restores the ``mm`` dispatch. The + non-detached path is unaffected either way (``should_fold`` folds early when the + weight requires grad), and the no-copy view keeps this shim free there too. + """ + + @staticmethod + def forward(ctx, x): + return x + + @staticmethod + def backward(ctx, grad): + if grad.is_contiguous(): + return grad.view(-1).view(grad.shape) + return grad.contiguous() + @runtime_checkable class DraftAdapter(Protocol): @@ -71,5 +101,6 @@ def project_mtp_hidden_to_logits(hidden_states_per_layer, model, detach_output_w logits_per_layer = [] for hidden in hidden_states_per_layer: logits, _ = model.output_layer(hidden, weight=output_weight) + logits = _CanonicalGradStrides.apply(logits) logits_per_layer.append(logits.transpose(0, 1).contiguous()) return logits_per_layer diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 73d958eec3..ed947bfa48 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -46,6 +46,9 @@ ) from skyrl.train.config import TrainerConfig +# One-shot guard for the MTP_PROFILE diagnostic capture (see forward_backward_mini_batch). +_MTP_PROFILE_DONE = False + class MegatronModelWrapper: def __init__( @@ -565,6 +568,8 @@ def loss_func(logits, data): return loss, metrics def forward_step(batch_iter, model): + if prof is not None: + prof.step() # NOTE(Charlie): despite the name, methods like `remove_left_padding()` are padding-agnostic # (can be left, or right) as it uses attention_mask to locate real tokens. Same thing # for recover_left_padding and setup_per_microbatch_replay_forward. Especially relevant @@ -679,15 +684,52 @@ def depad(tensor): # batch should be a list of micro-batches batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.actor_module)) - metrics_list = forward_backward_func( - forward_step_func=forward_step, - data_iterator=batch_generator, - model=self.actor_module, - num_microbatches=len(micro_batches), - seq_length=seq_len, - micro_batch_size=micro_batch_size, - forward_only=forward_only, - ) + # MTP_PROFILE=1: one-shot torch.profiler capture of steady-state microbatches (~9-13) of + # the first training mini-batch, for diagnosing the draft-path slowdown. Writes a per-rank + # kernel table to /tmp/mtp_prof_rank.txt (+ a chrome trace from rank 0). + prof = None + global _MTP_PROFILE_DONE + if os.environ.get("MTP_PROFILE") and not forward_only and not _MTP_PROFILE_DONE and len(micro_batches) >= 14: + _MTP_PROFILE_DONE = True + from torch.profiler import ProfilerActivity, profile, schedule + + _rank = torch.distributed.get_rank() + + def _dump_profile(p): + path = f"/tmp/mtp_prof_rank{_rank}.txt" + with open(path, "w") as f: + f.write( + p.key_averages(group_by_input_shape=True).table(sort_by="self_cuda_time_total", row_limit=40) + ) + f.write("\n\n===== by stack =====\n") + f.write(p.key_averages(group_by_stack_n=12).table(sort_by="self_cuda_time_total", row_limit=10)) + if _rank == 0: + p.export_chrome_trace("/tmp/mtp_prof_rank0_trace.json") + print(f"[MTP_PROFILE] wrote {path}", flush=True) + + prof = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=schedule(wait=8, warmup=2, active=3, repeat=1), + record_shapes=True, + with_stack=True, + on_trace_ready=_dump_profile, + ) + + if prof is not None: + prof.start() + try: + metrics_list = forward_backward_func( + forward_step_func=forward_step, + data_iterator=batch_generator, + model=self.actor_module, + num_microbatches=len(micro_batches), + seq_length=seq_len, + micro_batch_size=micro_batch_size, + forward_only=forward_only, + ) + finally: + if prof is not None: + prof.stop() # The decoupled MTP/draft loss is computed and logged per-microbatch inside loss_func # (metric key "mtp_loss"); no MTPLossLoggingHelper plumbing is needed. diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index b30e50c667..0975bc6ab3 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -782,6 +782,7 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: "UV_PYTHON", "UV_OFFLINE", "MTP_DEBUG", + "MTP_PROFILE", "PYTORCH_CUDA_ALLOC_CONF", # Debug/trace knobs — forwarded so they reach the worker actors, not just the driver. "CUDA_LAUNCH_BLOCKING", diff --git a/tests/backends/skyrl_train/mtp/test_adapter.py b/tests/backends/skyrl_train/mtp/test_adapter.py new file mode 100644 index 0000000000..d364f0d6c3 --- /dev/null +++ b/tests/backends/skyrl_train/mtp/test_adapter.py @@ -0,0 +1,78 @@ +"""CPU unit tests for the draft-projection adapter helpers. + +Covers ``_CanonicalGradStrides``: the identity shim that re-canonicalizes gradient +strides between the (frozen-weight) output projection and the adapter's transpose. +Without it, a micro-batch-1 transpose-backward grad keeps a stale stride in its +size-1 batch dim; megatron's ``LinearWithFrozenWeight.backward`` then fails +``torch.matmul``'s fold-to-mm stride check (which does not skip size-1 dims) and +dispatches a batch=seq, M=1 broadcast bmm that is ~100x slower than the flat mm. + +uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_adapter.py +""" + +import torch + +from skyrl.backends.skyrl_train.mtp.adapter import _CanonicalGradStrides + + +def _passes_matmul_fold_check(t: torch.Tensor) -> bool: + """Mirror torch's ``should_fold`` stride loop: foldable iff the leading dims are + dense in order, with no size-1 skipping (the rule the real dispatch applies).""" + sizes, strides = t.shape, t.stride() + return all(strides[i] == strides[i + 1] * sizes[i + 1] for i in range(t.dim() - 2)) + + +def _projection_chain(x, use_shim: bool): + """Mimic project_mtp_hidden_to_logits' post-projection ops on a [seq, 1, vocab] + tensor: (shim) -> transpose(0,1) -> contiguous (a no-op view for batch 1).""" + y = _CanonicalGradStrides.apply(x) if use_shim else x + return y.transpose(0, 1).contiguous() + + +def test_canonical_grad_strides_restores_foldable_grad(): + # Capture the raw gradient flowing INTO the projection via a tensor hook — a + # mid-graph consumer (like LinearWithFrozenWeight.backward) sees exactly this + # tensor. Leaf `.grad` would not do: AccumulateGrad re-lays-out the grad to + # match the leaf's strides, hiding the stale-stride problem. + S, V = 6, 10 + base = torch.randn(S, 1, V, requires_grad=True) + + def run(use_shim): + captured = [] + x = base.clone() # mid-graph tensor standing in for the projection output + x.register_hook(captured.append) + out = _projection_chain(x, use_shim=use_shim) + # Freshly-allocated canonical grad, like the real upstream (the depad's + # zeros+scatter). ones_like(out) would inherit the view's strides and + # accidentally survive the transpose with a foldable layout. + out.backward(torch.ones(out.shape)) + return captured[0] + + grad_no_shim = run(use_shim=False) + grad_shim = run(use_shim=True) + + # The shim must not change gradient values, only the stride layout. + assert torch.equal(grad_no_shim, grad_shim) + # Without the shim the transpose-backward grad carries a stale stride in the + # size-1 dim and would fall off matmul's mm fast path; with it, it folds. + assert not _passes_matmul_fold_check(grad_no_shim) + assert _passes_matmul_fold_check(grad_shim) + + +def test_canonical_grad_strides_identity_forward(): + x = torch.randn(4, 1, 8, requires_grad=True) + y = _CanonicalGradStrides.apply(x) + assert y.data_ptr() == x.data_ptr() + assert torch.equal(y, x) + + +def test_canonical_grad_strides_handles_noncontiguous_grad(): + # A genuinely non-contiguous upstream grad (not just stale size-1 strides) must + # come out contiguous and value-identical. + x = torch.randn(5, 3, 4, requires_grad=True) + y = _CanonicalGradStrides.apply(x) + g = torch.randn(4, 3, 5).permute(2, 1, 0) # non-contiguous grad + assert not g.is_contiguous() + y.backward(g) + assert x.grad.is_contiguous() + assert torch.equal(x.grad, g) From cad26fc0bd399bec9e030903e96dbe63915f367d Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Mon, 15 Jun 2026 19:08:01 +0000 Subject: [PATCH 14/28] fix vllm metrics logger throughput calculation using step time --- skyrl/train/fully_async_trainer.py | 2 + skyrl/train/trainer.py | 8 +- skyrl/train/utils/vllm_metrics_scraper.py | 20 +++-- tests/train/test_vllm_metrics_scraper.py | 105 ++++++++++++++++++++++ 4 files changed, 129 insertions(+), 6 deletions(-) diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index fe31b2fe55..bee8773518 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -459,6 +459,8 @@ async def train(self): timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()} if self._vllm_metrics_scraper is not None: + # Generation overlaps training here, so throughput falls + # back to the wall-clock interval. timing_payload.update(await self._vllm_metrics_scraper.sample()) self.tracker.log(timing_payload, step=self.global_step) self.all_timings = {} diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index cf75bed849..851b21b0ba 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -437,7 +437,13 @@ async def train(self): **{f"timing/{k}": v for k, v in self.all_timings.items()}, } if self._vllm_metrics_scraper is not None: - log_payload.update(await self._vllm_metrics_scraper.sample()) + # Engine only generates during the "generate" phase, so + # throughput divides by that, not the full step time. + log_payload.update( + await self._vllm_metrics_scraper.sample( + generation_time_s=self.all_timings.get("generate") + ) + ) if self._ray_gpu_monitor is not None: log_payload.update(self._ray_gpu_monitor.flush()) diff --git a/skyrl/train/utils/vllm_metrics_scraper.py b/skyrl/train/utils/vllm_metrics_scraper.py index d3099dd718..770a912ca4 100644 --- a/skyrl/train/utils/vllm_metrics_scraper.py +++ b/skyrl/train/utils/vllm_metrics_scraper.py @@ -194,11 +194,16 @@ async def _fetch_all(self) -> ParsedSamples: merged[key] = value return merged - async def sample(self) -> Dict[str, float]: + async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, float]: """Return a dict of ``vllm/...`` scalars for the current step. Empty if no agents are reachable or if no vLLM samples are present yet (e.g. before any inference has run). + + ``generation_time_s`` is the time the engine spent generating since the + previous ``sample()`` call; it is the denominator for the throughput + metrics. Pass ``None`` when generation runs for the whole interval + (e.g. fully-async overlap) to fall back to the wall-clock interval. """ if not self._urls: return {} @@ -230,14 +235,19 @@ async def sample(self) -> Dict[str, float]: # Derived metrics need a previous snapshot to take deltas. if self._prev_aggregated is not None and self._prev_timestamp is not None: dt = max(now - self._prev_timestamp, 1e-9) - out.update(self._derive(snapshot, self._prev_aggregated, dt)) + # Throughput is per generation-second, not per step-second. + if generation_time_s is not None and generation_time_s > 0: + throughput_window_s = generation_time_s + else: + throughput_window_s = dt + out.update(self._derive(snapshot, self._prev_aggregated, throughput_window_s)) self._prev_aggregated = snapshot self._prev_timestamp = now return out @staticmethod - def _derive(cur: Dict[str, float], prev: Dict[str, float], dt: float) -> Dict[str, float]: + def _derive(cur: Dict[str, float], prev: Dict[str, float], throughput_window_s: float) -> Dict[str, float]: out: Dict[str, float] = {} def delta(name: str) -> Optional[float]: @@ -249,11 +259,11 @@ def delta(name: str) -> Optional[float]: gen_d = delta(_COUNTER_GENERATION_TOKENS) if gen_d is not None: - out["vllm/generation_throughput_tok_s"] = gen_d / dt + out["vllm/generation_throughput_tok_s"] = gen_d / throughput_window_s prompt_d = delta(_COUNTER_PROMPT_TOKENS) if prompt_d is not None: - out["vllm/prompt_throughput_tok_s"] = prompt_d / dt + out["vllm/prompt_throughput_tok_s"] = prompt_d / throughput_window_s q_d = delta(_COUNTER_PREFIX_QUERIES) h_d = delta(_COUNTER_PREFIX_HITS) diff --git a/tests/train/test_vllm_metrics_scraper.py b/tests/train/test_vllm_metrics_scraper.py index bb6d67d378..9066ca4d0f 100644 --- a/tests/train/test_vllm_metrics_scraper.py +++ b/tests/train/test_vllm_metrics_scraper.py @@ -218,6 +218,111 @@ async def fake_fetch_all(): assert out["vllm/kv_cache_usage_perc"] == pytest.approx(0.35) +@pytest.mark.asyncio +async def test_scraper_throughput_uses_generation_time_not_wall_clock(): + """Throughput divides token deltas by generation time, not the full step. + + The wall-clock gap between samples is 10s, but only 2s was spent + generating. Throughput must use the 2s window; latency/hit-rate metrics + (which don't depend on time) are unaffected. + """ + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + snap1 = _snapshot( + running=4, + waiting=2, + kv=0.5, + prefix_q=100, + prefix_h=80, + prompt_toks=1000, + gen_toks=500, + ttft_sum=2.0, + ttft_count=10, + itl_sum=1.0, + itl_count=200, + ) + # +250 gen tokens, +100 prompt tokens over a 10s wall-clock gap of which + # only 2s was generation. + snap2 = _snapshot( + running=5, + waiting=1, + kv=0.7, + prefix_q=150, + prefix_h=120, + prompt_toks=1100, + gen_toks=750, + ttft_sum=3.0, + ttft_count=15, + itl_sum=1.5, + itl_count=300, + ) + + texts = iter([snap1, snap2]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + times = iter([1000.0, 1010.0]) # 10s wall-clock apart + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + await scraper.sample() + out = await scraper.sample(generation_time_s=2.0) + + # 250 tokens / 2s generation == 125 tok/s (NOT 250/10 == 25). + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(125.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + # Time-independent derived metrics are unchanged by the window choice. + assert out["vllm/prefix_cache_hit_rate"] == pytest.approx(40.0 / 50.0) + assert out["vllm/ttft_seconds_avg"] == pytest.approx(1.0 / 5.0) + assert out["vllm/tpot_seconds_avg"] == pytest.approx(0.5 / 100.0) + + +@pytest.mark.asyncio +async def test_scraper_throughput_falls_back_to_wall_clock_without_gen_time(): + """With no generation_time_s (and with non-positive values), use dt.""" + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + common = dict( + running=0, + waiting=0, + kv=0.0, + prefix_q=0, + prefix_h=0, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + snap1 = _snapshot(prompt_toks=0, gen_toks=0, **common) + snap2 = _snapshot(prompt_toks=100, gen_toks=200, **common) + texts = iter([snap1, snap2]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + times = iter([1000.0, 1002.0]) # 2s wall-clock apart + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + await scraper.sample() + # generation_time_s=0 is non-positive -> fall back to wall-clock dt. + out = await scraper.sample(generation_time_s=0.0) + + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(100.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + + @pytest.mark.asyncio async def test_scraper_handles_counter_reset(): scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) From 8e6b355c84954a7d8bc0674dc37adbbf7186a117 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Tue, 16 Jun 2026 01:34:21 +0000 Subject: [PATCH 15/28] fix logging --- .../megatron/run_megatron_dapo_qwen3.5_9b.sh | 4 +- ...run_megatron_dapo_qwen3.5_9b_specdecode.sh | 6 +- .../vllm/spec_decode_metrics.py | 18 +- skyrl/train/evaluate.py | 12 ++ skyrl/train/fully_async_trainer.py | 2 + skyrl/train/trainer.py | 33 ++- skyrl/train/utils/vllm_metrics_scraper.py | 151 +++++++++++--- .../mtp/test_spec_decode_metrics.py | 30 +++ tests/train/test_vllm_metrics_scraper.py | 191 ++++++++++++++++++ 9 files changed, 401 insertions(+), 46 deletions(-) diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh index 042ef75985..30f0531229 100644 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b.sh @@ -129,7 +129,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ trainer.micro_forward_batch_size_per_gpu=1 \ trainer.micro_train_batch_size_per_gpu=1 \ - trainer.ckpt_interval=-1 \ + trainer.ckpt_interval=50 \ trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ trainer.policy.optimizer_config.lr=$LR \ @@ -146,7 +146,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ generator.inference_engine.gpu_memory_utilization=0.5 \ trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo" \ + trainer.project_name="qwen3_5_dapo_2" \ trainer.run_name="nosd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.export_path="/mnt/local_storage/exports/dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.hf_save_interval=300 \ diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh index cb19c3f6ea..bc134045ce 100644 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh @@ -101,7 +101,7 @@ TIS_TYPE=token MTP_ENABLED=true MTP_NUM_SPECULATIVE_TOKENS=3 MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) -MTP_LOSS_WEIGHT=0.5 +MTP_LOSS_WEIGHT=1.0 # NOTE: trainer.policy.megatron_config.mtp_detach_shared_output now defaults to true: the draft loss # trains ONLY the MTP-head params; the tied embedding/lm_head is detached (output projection AND the # MTP block's re-embedding), so the draft gradient no longer nudges the policy's own logits. Set it @@ -167,7 +167,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ trainer.micro_forward_batch_size_per_gpu=1 \ trainer.micro_train_batch_size_per_gpu=1 \ - trainer.ckpt_interval=-1 \ + trainer.ckpt_interval=50 \ trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ trainer.policy.optimizer_config.lr=$LR \ @@ -189,7 +189,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo" \ + trainer.project_name="qwen3_5_dapo_2" \ trainer.run_name="sd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.export_path="/mnt/local_storage/exports/sd_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.hf_save_interval=300 \ diff --git a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py index 1295f87d61..77be1263cd 100644 --- a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py +++ b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py @@ -87,17 +87,21 @@ def merge_spec_decode_counters(totals: dict, stats: dict) -> None: totals[key] = totals.get(key, 0) + int(value) -def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> tuple[dict, Optional[dict]]: +def acceptance_rate_metrics( + cumulative: Optional[dict], prev: Optional[dict], prefix: str = "vllm/" +) -> tuple[dict, Optional[dict]]: """Turn cumulative spec-decode counters into per-step (delta) metrics. Args: cumulative: counters read this step (from the engines), or None if speculative decoding is disabled / unsupported. prev: the ``cumulative`` from the previous step (None on the first step). + prefix: metric-key prefix. ``"vllm/"`` for the train rollout; pass ``"vllm/eval/"`` to + attribute the eval rollout's draft/accept delta separately. Returns: - ``(metrics, new_prev)`` where ``metrics`` has ``vllm/draft_*`` keys (empty when there are no - stats) and ``new_prev`` is the snapshot to pass back next step. + ``(metrics, new_prev)`` where ``metrics`` has ``{prefix}draft_*`` keys (empty when there are + no stats) and ``new_prev`` is the snapshot to pass back next step. """ if not cumulative: return {}, prev @@ -106,12 +110,12 @@ def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> drafted = cumulative.get("num_draft_tokens", 0) - prev.get("num_draft_tokens", 0) accepted = cumulative.get("num_accepted_tokens", 0) - prev.get("num_accepted_tokens", 0) metrics: dict[str, Any] = { - "vllm/draft_num_draft_tokens": drafted, - "vllm/draft_num_accepted_tokens": accepted, + f"{prefix}draft_num_draft_tokens": drafted, + f"{prefix}draft_num_accepted_tokens": accepted, } if drafted > 0: # Acceptance rate = accepted draft tokens / total drafted tokens this step. - metrics["vllm/draft_acceptance_rate"] = accepted / drafted + metrics[f"{prefix}draft_acceptance_rate"] = accepted / drafted # Per-position rates: fraction of draft rounds this step whose k-th speculated token was # accepted (same definition vLLM uses in its own per-position logging). Position keys are # 1-based: pos_1 = first drafted token. Acceptance halts at the first rejection, so the rates @@ -121,5 +125,5 @@ def acceptance_rate_metrics(cumulative: Optional[dict], prev: Optional[dict]) -> if drafts > 0: for i, n in enumerate(per_pos): prev_n = int(prev_per_pos[i]) if i < len(prev_per_pos) else 0 - metrics[f"vllm/draft_acceptance_rate_pos_{i + 1}"] = (int(n) - prev_n) / drafts + metrics[f"{prefix}draft_acceptance_rate_pos_{i + 1}"] = (int(n) - prev_n) / drafts return metrics, cumulative diff --git a/skyrl/train/evaluate.py b/skyrl/train/evaluate.py index f24c84f2bd..5917ee8499 100644 --- a/skyrl/train/evaluate.py +++ b/skyrl/train/evaluate.py @@ -1,3 +1,4 @@ +import time from collections import defaultdict from pathlib import Path from typing import Any, Dict, List @@ -58,6 +59,7 @@ async def evaluate( concat_env_extras: List[Dict[str, Any]] = [] concat_uids: List[str] = [] sampling_params = cfg.generator.eval_sampling_params + eval_generate_time = 0.0 pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress") for _, prompts in enumerate(eval_dataloader): pbar.update(1) @@ -69,7 +71,9 @@ async def evaluate( "eval", global_step, ) + gen_start = time.monotonic() generator_output: GeneratorOutput = await generator.generate(generator_input) + eval_generate_time += time.monotonic() - gen_start validate_generator_output(len(generator_input["prompts"]), generator_output) generator_outputs.append(generator_output) concat_all_envs.extend(generator_input["env_classes"]) @@ -127,6 +131,9 @@ async def evaluate( eval_metrics, ) + # Rollout-only time, used as the vLLM throughput denominator on eval steps. + eval_metrics["timing/eval_generate"] = eval_generate_time + return eval_metrics @@ -160,6 +167,7 @@ async def evaluate_step_wise( concat_env_extras: List[Dict[str, Any]] = [] concat_uids: List[str] = [] sampling_params = cfg.generator.eval_sampling_params + eval_generate_time = 0.0 pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress") for _, prompts in enumerate(eval_dataloader): pbar.update(1) @@ -171,7 +179,9 @@ async def evaluate_step_wise( "eval", global_step, ) + gen_start = time.monotonic() generator_output: GeneratorOutput = await generator.generate(generator_input) + eval_generate_time += time.monotonic() - gen_start traj_id_to_input = { traj_id.instance_id: {"env_class": env_class, "env_extras": env_extra} for traj_id, env_class, env_extra in zip( @@ -244,4 +254,6 @@ async def evaluate_step_wise( eval_metrics, ) + eval_metrics["timing/eval_generate"] = eval_generate_time + return eval_metrics diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index fe31b2fe55..bee8773518 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -459,6 +459,8 @@ async def train(self): timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()} if self._vllm_metrics_scraper is not None: + # Generation overlaps training here, so throughput falls + # back to the wall-clock interval. timing_payload.update(await self._vllm_metrics_scraper.sample()) self.tracker.log(timing_payload, step=self.global_step) self.all_timings = {} diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index d7913fa53e..5da4128c9a 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -429,10 +429,18 @@ async def train(self): or self.global_step == self.total_training_steps ) if force_eval or interval_eval: + # Mark the vLLM counters before eval so the eval rollout is + # attributed to vllm/eval/* rather than blended into the + # train rollout's vllm/* throughput. + if self._vllm_metrics_scraper is not None: + await self._vllm_metrics_scraper.mark_pre_eval() self._fire("on_eval_start") with Timer("eval", self.all_timings): eval_metrics = await self.eval() self.all_metrics.update(eval_metrics) + # Attribute the eval rollout's draft acceptance to vllm/eval/* and advance the + # spec-decode baseline so eval counts don't leak into the next train step. + await self._record_spec_decode_metrics(prefix="vllm/eval/") self._fire("on_eval_end", metrics=eval_metrics) log_payload = { @@ -440,7 +448,16 @@ async def train(self): **{f"timing/{k}": v for k, v in self.all_timings.items()}, } if self._vllm_metrics_scraper is not None: - log_payload.update(await self._vllm_metrics_scraper.sample()) + # vllm/* covers the train rollout (divided by generate time); + # on eval steps vllm/eval/* covers the eval rollout (divided + # by the eval rollout time). Each rollout's tokens are thus + # divided by its own time, avoiding the eval-step spike. + log_payload.update( + await self._vllm_metrics_scraper.sample_split( + generate_time_s=self.all_timings.get("generate"), + eval_generate_time_s=self.all_metrics.get("timing/eval_generate"), + ) + ) if self._ray_gpu_monitor is not None: log_payload.update(self._ray_gpu_monitor.flush()) @@ -841,14 +858,20 @@ def convert_to_training_input(self, generator_output: GeneratorOutput, uids: Lis return training_input @torch.no_grad() - async def _record_spec_decode_metrics(self) -> None: - """Record the vLLM speculative-decoding (MTP draft) acceptance rate for this rollout step. + async def _record_spec_decode_metrics(self, prefix: str = "vllm/") -> None: + """Record the vLLM speculative-decoding (MTP draft) acceptance rate for this rollout. Reads the engines' cumulative draft/accept counters and logs the per-step delta as - ``vllm/draft_acceptance_rate`` (+ raw counts), plus one ``vllm/draft_acceptance_rate_pos_k`` + ``{prefix}draft_acceptance_rate`` (+ raw counts), plus one ``{prefix}draft_acceptance_rate_pos_k`` per draft position when num_speculative_tokens > 1 (per-depth acceptance decay). No-op when speculative decoding is disabled or the backend doesn't expose the stats. Best-effort: never fails the training step. + + ``prefix`` is ``"vllm/"`` for the train rollout. Because the cumulative counters keep + ``self._prev_spec_decode`` at the post-train-rollout value (nothing generates between the + train rollout and eval), calling this again with ``prefix="vllm/eval/"`` right after the + eval rollout reports the eval delta under ``vllm/eval/*`` AND advances the baseline so the + eval counts don't leak into the next train step's acceptance rate. """ from skyrl.backends.skyrl_train.inference_engines.vllm.spec_decode_metrics import ( acceptance_rate_metrics, @@ -859,7 +882,7 @@ async def _record_spec_decode_metrics(self) -> None: except Exception as e: logger.warning(f"Failed to read vLLM spec-decode metrics: {e}") return - metrics, self._prev_spec_decode = acceptance_rate_metrics(cumulative, self._prev_spec_decode) + metrics, self._prev_spec_decode = acceptance_rate_metrics(cumulative, self._prev_spec_decode, prefix=prefix) self.all_metrics.update(metrics) async def generate( diff --git a/skyrl/train/utils/vllm_metrics_scraper.py b/skyrl/train/utils/vllm_metrics_scraper.py index d3099dd718..4a31ad357b 100644 --- a/skyrl/train/utils/vllm_metrics_scraper.py +++ b/skyrl/train/utils/vllm_metrics_scraper.py @@ -140,9 +140,17 @@ def discover_ray_metrics_urls() -> List[str]: class VLLMMetricsScraper: """Per-step snapshot of selected vLLM metrics from Ray's metrics agents. - The first ``sample()`` call establishes a baseline; rate-style metrics - (throughput, hit rate, average latency) are reported starting from the - second call. + The first sample establishes a baseline; rate-style metrics (throughput, + hit rate, average latency) are reported starting from the second call. + + Two usage modes share the same plumbing: + + * ``sample()`` — one snapshot per step, deltas vs. the previous step. Used + by the fully-async trainer, where generation runs for the whole interval. + * ``mark_pre_eval()`` + ``sample_split()`` — the synchronous trainer marks + the counters right before the eval rollout so the train rollout is + reported under ``vllm/*`` and the (much larger, differently-batched) eval + rollout under ``vllm/eval/*`` instead of being blended into one number. """ def __init__( @@ -154,6 +162,8 @@ def __init__( self._timeout = request_timeout_s self._prev_aggregated: Optional[Dict[str, float]] = None self._prev_timestamp: Optional[float] = None + # Counter snapshot taken right before an eval rollout (sync trainer). + self._mid_snapshot: Optional[Dict[str, float]] = None self._client: Optional[httpx.AsyncClient] = None self._warned_empty = False if not self._urls: @@ -194,14 +204,14 @@ async def _fetch_all(self) -> ParsedSamples: merged[key] = value return merged - async def sample(self) -> Dict[str, float]: - """Return a dict of ``vllm/...`` scalars for the current step. + async def _read_snapshot(self) -> Optional[Dict[str, float]]: + """Scrape every agent and reduce to one cumulative value per metric. - Empty if no agents are reachable or if no vLLM samples are present - yet (e.g. before any inference has run). + Returns ``None`` when no endpoints are configured. Counters are summed + across replicas; gauges are averaged. """ if not self._urls: - return {} + return None parsed = await self._fetch_all() if not parsed and not self._warned_empty: @@ -212,32 +222,111 @@ async def sample(self) -> Dict[str, float]: ) self._warned_empty = True - now = time.monotonic() - sums = aggregate(parsed, _SUM_METRICS, how="sum") means = aggregate(parsed, _MEAN_METRICS, how="mean") - snapshot = {**sums, **means} + return {**sums, **means} - out: Dict[str, float] = {} - # Instantaneous gauges expose directly. - if _GAUGE_NUM_RUNNING in snapshot: - out["vllm/num_requests_running"] = snapshot[_GAUGE_NUM_RUNNING] - if _GAUGE_NUM_WAITING in snapshot: - out["vllm/num_requests_waiting"] = snapshot[_GAUGE_NUM_WAITING] - if _GAUGE_KV_CACHE_USAGE in snapshot: - out["vllm/kv_cache_usage_perc"] = snapshot[_GAUGE_KV_CACHE_USAGE] + async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, float]: + """Return a dict of ``vllm/...`` scalars for the current step. - # Derived metrics need a previous snapshot to take deltas. + Empty if no agents are reachable or if no vLLM samples are present + yet (e.g. before any inference has run). + + ``generation_time_s`` is the time the engine spent generating since the + previous ``sample()`` call; it is the denominator for the throughput + metrics. Pass ``None`` when generation runs for the whole interval + (e.g. fully-async overlap) to fall back to the wall-clock interval. + """ + snapshot = await self._read_snapshot() + if snapshot is None: + return {} + + now = time.monotonic() if self._prev_aggregated is not None and self._prev_timestamp is not None: dt = max(now - self._prev_timestamp, 1e-9) - out.update(self._derive(snapshot, self._prev_aggregated, dt)) + # Throughput is per generation-second, not per step-second. + window = generation_time_s if (generation_time_s is not None and generation_time_s > 0) else dt + out = self._window_metrics(self._prev_aggregated, snapshot, window, "vllm/") + else: + # First call: gauges only, no deltas yet. + out = self._window_metrics(None, snapshot, None, "vllm/") self._prev_aggregated = snapshot self._prev_timestamp = now return out + async def mark_pre_eval(self) -> None: + """Snapshot the counters right before the eval rollout. + + Lets :meth:`sample_split` attribute eval-rollout generation to + ``vllm/eval/*`` separately from the train rollout. No-op effect if no + endpoints are configured (the stored mark is simply ``None``). + """ + self._mid_snapshot = await self._read_snapshot() + + async def sample_split( + self, + *, + generate_time_s: Optional[float] = None, + eval_generate_time_s: Optional[float] = None, + ) -> Dict[str, float]: + """Synchronous-trainer sampling that separates train and eval rollouts. + + Emits ``vllm/*`` for the train rollout (delta from the previous step up + to the pre-eval mark, divided by ``generate_time_s``). When + :meth:`mark_pre_eval` was called this step, also emits ``vllm/eval/*`` + for the eval rollout (delta from the mark to now, divided by + ``eval_generate_time_s``). + """ + cur = await self._read_snapshot() + mid = self._mid_snapshot + self._mid_snapshot = None + if cur is None: + return {} + + # Train rollout ends at the pre-eval mark when eval ran this step, + # otherwise nothing generated after it, so the current snapshot is fine. + train_end = mid if mid is not None else cur + train_window = generate_time_s if (generate_time_s is not None and generate_time_s > 0) else None + out = self._window_metrics(self._prev_aggregated, train_end, train_window, "vllm/") + + if mid is not None: + eval_window = ( + eval_generate_time_s if (eval_generate_time_s is not None and eval_generate_time_s > 0) else None + ) + out.update(self._window_metrics(mid, cur, eval_window, "vllm/eval/")) + + self._prev_aggregated = cur + return out + + @classmethod + def _window_metrics( + cls, + prev: Optional[Dict[str, float]], + cur: Optional[Dict[str, float]], + throughput_window_s: Optional[float], + prefix: str, + ) -> Dict[str, float]: + """Gauges from ``cur`` plus derived rates over ``cur - prev``.""" + if cur is None: + return {} + out: Dict[str, float] = {} + # Instantaneous gauges expose directly. + if _GAUGE_NUM_RUNNING in cur: + out[f"{prefix}num_requests_running"] = cur[_GAUGE_NUM_RUNNING] + if _GAUGE_NUM_WAITING in cur: + out[f"{prefix}num_requests_waiting"] = cur[_GAUGE_NUM_WAITING] + if _GAUGE_KV_CACHE_USAGE in cur: + out[f"{prefix}kv_cache_usage_perc"] = cur[_GAUGE_KV_CACHE_USAGE] + # Derived metrics need a previous snapshot to take deltas. + if prev is not None: + out.update(cls._derive(cur, prev, throughput_window_s, prefix)) + return out + @staticmethod - def _derive(cur: Dict[str, float], prev: Dict[str, float], dt: float) -> Dict[str, float]: + def _derive( + cur: Dict[str, float], prev: Dict[str, float], throughput_window_s: Optional[float], prefix: str + ) -> Dict[str, float]: out: Dict[str, float] = {} def delta(name: str) -> Optional[float]: @@ -247,27 +336,31 @@ def delta(name: str) -> Optional[float]: # Counter resets (engine restart) shouldn't crash; just skip. return d if d >= 0 else None + # Throughput needs a positive time window; the rate/latency metrics + # below are time-independent, so they emit regardless. + has_window = throughput_window_s is not None and throughput_window_s > 0 + gen_d = delta(_COUNTER_GENERATION_TOKENS) - if gen_d is not None: - out["vllm/generation_throughput_tok_s"] = gen_d / dt + if gen_d is not None and has_window: + out[f"{prefix}generation_throughput_tok_s"] = gen_d / throughput_window_s prompt_d = delta(_COUNTER_PROMPT_TOKENS) - if prompt_d is not None: - out["vllm/prompt_throughput_tok_s"] = prompt_d / dt + if prompt_d is not None and has_window: + out[f"{prefix}prompt_throughput_tok_s"] = prompt_d / throughput_window_s q_d = delta(_COUNTER_PREFIX_QUERIES) h_d = delta(_COUNTER_PREFIX_HITS) if q_d is not None and h_d is not None and q_d > 0: - out["vllm/prefix_cache_hit_rate"] = h_d / q_d + out[f"{prefix}prefix_cache_hit_rate"] = h_d / q_d ttft_sum_d = delta(_HIST_TTFT_SUM) ttft_count_d = delta(_HIST_TTFT_COUNT) if ttft_sum_d is not None and ttft_count_d is not None and ttft_count_d > 0: - out["vllm/ttft_seconds_avg"] = ttft_sum_d / ttft_count_d + out[f"{prefix}ttft_seconds_avg"] = ttft_sum_d / ttft_count_d itl_sum_d = delta(_HIST_ITL_SUM) itl_count_d = delta(_HIST_ITL_COUNT) if itl_sum_d is not None and itl_count_d is not None and itl_count_d > 0: - out["vllm/tpot_seconds_avg"] = itl_sum_d / itl_count_d + out[f"{prefix}tpot_seconds_avg"] = itl_sum_d / itl_count_d return out diff --git a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py index 012602558a..da8546784e 100644 --- a/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py +++ b/tests/backends/skyrl_train/mtp/test_spec_decode_metrics.py @@ -160,6 +160,36 @@ def test_acceptance_rate_per_position_depth_one_single_key(): assert abs(metrics["vllm/draft_acceptance_rate_pos_1"] - 0.75) < 1e-9 +def test_acceptance_rate_prefix_separates_train_and_eval(): + # Mirrors the trainer's eval-step sequence: a train rollout records under vllm/*, then the eval + # rollout records under vllm/eval/* using the SAME advancing baseline, so eval counts don't leak + # into the next train step. + prev = None + # Train rollout: cumulative after train generate. + train, prev = acceptance_rate_metrics( + {"num_drafts": 10, "num_draft_tokens": 20, "num_accepted_tokens": 12}, prev, prefix="vllm/" + ) + assert train["vllm/draft_num_draft_tokens"] == 20 + assert abs(train["vllm/draft_acceptance_rate"] - 0.6) < 1e-9 + + # Eval rollout: cumulative grows a lot; reported under vllm/eval/* as the delta vs the train + # baseline (drafted 120-20=100, accepted 90-12=78 -> 0.78). + eval_m, prev = acceptance_rate_metrics( + {"num_drafts": 60, "num_draft_tokens": 120, "num_accepted_tokens": 90}, prev, prefix="vllm/eval/" + ) + assert eval_m["vllm/eval/draft_num_draft_tokens"] == 100 + assert abs(eval_m["vllm/eval/draft_acceptance_rate"] - 0.78) < 1e-9 + assert not any(k.startswith("vllm/draft_") for k in eval_m) + + # Next train step: delta is from the post-eval baseline, so eval is NOT counted again + # (drafted 140-120=20, accepted 105-90=15 -> 0.75). + train2, prev = acceptance_rate_metrics( + {"num_drafts": 70, "num_draft_tokens": 140, "num_accepted_tokens": 105}, prev, prefix="vllm/" + ) + assert train2["vllm/draft_num_draft_tokens"] == 20 + assert abs(train2["vllm/draft_acceptance_rate"] - 0.75) < 1e-9 + + def test_acceptance_rate_no_spec_decode_returns_empty(): metrics, prev = acceptance_rate_metrics(None, None) assert metrics == {} diff --git a/tests/train/test_vllm_metrics_scraper.py b/tests/train/test_vllm_metrics_scraper.py index bb6d67d378..48f4e7b577 100644 --- a/tests/train/test_vllm_metrics_scraper.py +++ b/tests/train/test_vllm_metrics_scraper.py @@ -218,6 +218,111 @@ async def fake_fetch_all(): assert out["vllm/kv_cache_usage_perc"] == pytest.approx(0.35) +@pytest.mark.asyncio +async def test_scraper_throughput_uses_generation_time_not_wall_clock(): + """Throughput divides token deltas by generation time, not the full step. + + The wall-clock gap between samples is 10s, but only 2s was spent + generating. Throughput must use the 2s window; latency/hit-rate metrics + (which don't depend on time) are unaffected. + """ + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + snap1 = _snapshot( + running=4, + waiting=2, + kv=0.5, + prefix_q=100, + prefix_h=80, + prompt_toks=1000, + gen_toks=500, + ttft_sum=2.0, + ttft_count=10, + itl_sum=1.0, + itl_count=200, + ) + # +250 gen tokens, +100 prompt tokens over a 10s wall-clock gap of which + # only 2s was generation. + snap2 = _snapshot( + running=5, + waiting=1, + kv=0.7, + prefix_q=150, + prefix_h=120, + prompt_toks=1100, + gen_toks=750, + ttft_sum=3.0, + ttft_count=15, + itl_sum=1.5, + itl_count=300, + ) + + texts = iter([snap1, snap2]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + times = iter([1000.0, 1010.0]) # 10s wall-clock apart + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + await scraper.sample() + out = await scraper.sample(generation_time_s=2.0) + + # 250 tokens / 2s generation == 125 tok/s (NOT 250/10 == 25). + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(125.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + # Time-independent derived metrics are unchanged by the window choice. + assert out["vllm/prefix_cache_hit_rate"] == pytest.approx(40.0 / 50.0) + assert out["vllm/ttft_seconds_avg"] == pytest.approx(1.0 / 5.0) + assert out["vllm/tpot_seconds_avg"] == pytest.approx(0.5 / 100.0) + + +@pytest.mark.asyncio +async def test_scraper_throughput_falls_back_to_wall_clock_without_gen_time(): + """With no generation_time_s (and with non-positive values), use dt.""" + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + common = dict( + running=0, + waiting=0, + kv=0.0, + prefix_q=0, + prefix_h=0, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + snap1 = _snapshot(prompt_toks=0, gen_toks=0, **common) + snap2 = _snapshot(prompt_toks=100, gen_toks=200, **common) + texts = iter([snap1, snap2]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + times = iter([1000.0, 1002.0]) # 2s wall-clock apart + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + await scraper.sample() + # generation_time_s=0 is non-positive -> fall back to wall-clock dt. + out = await scraper.sample(generation_time_s=0.0) + + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(100.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + + @pytest.mark.asyncio async def test_scraper_handles_counter_reset(): scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) @@ -267,6 +372,92 @@ def test_scraper_with_no_urls_is_noop(): assert out == {} +@pytest.mark.asyncio +async def test_sample_split_separates_train_and_eval_rollouts(): + """vllm/* uses the train window; vllm/eval/* uses the eval window. + + Timeline of one eval step: + baseline gen=500 prompt=1000 + pre-eval gen=600 prompt=1100 (train rollout: +100 gen, +100 prompt) + log-time gen=1100 prompt=1300 (eval rollout: +500 gen, +200 prompt) + with generate_time=2s and eval_generate_time=5s. + """ + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + def snap(gen, prompt): + return _snapshot( + running=1, + waiting=0, + kv=0.4, + prefix_q=0, + prefix_h=0, + prompt_toks=prompt, + gen_toks=gen, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + + baseline = snap(500, 1000) + pre_eval = snap(600, 1100) + log_time = snap(1100, 1300) + texts = iter([baseline, pre_eval, log_time]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + with patch.object(scraper, "_fetch_all", fake_fetch_all): + # Step before eval: establishes the baseline (no throughput yet). + first = await scraper.sample_split(generate_time_s=1.0) + assert "vllm/generation_throughput_tok_s" not in first + + # Eval step: mark before eval, then sample after. + await scraper.mark_pre_eval() + out = await scraper.sample_split(generate_time_s=2.0, eval_generate_time_s=5.0) + + # Train rollout: (600-500)/2 == 50, (1100-1000)/2 == 50. + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(50.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + # Eval rollout: (1100-600)/5 == 100, (1300-1100)/5 == 40. + assert out["vllm/eval/generation_throughput_tok_s"] == pytest.approx(100.0) + assert out["vllm/eval/prompt_throughput_tok_s"] == pytest.approx(40.0) + + +@pytest.mark.asyncio +async def test_sample_split_without_eval_emits_only_train_metrics(): + """On a non-eval step (no mark_pre_eval), no vllm/eval/* keys appear.""" + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + def snap(gen, prompt): + return _snapshot( + running=1, + waiting=0, + kv=0.4, + prefix_q=0, + prefix_h=0, + prompt_toks=prompt, + gen_toks=gen, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + + texts = iter([snap(0, 0), snap(200, 100)]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + with patch.object(scraper, "_fetch_all", fake_fetch_all): + await scraper.sample_split(generate_time_s=1.0) + out = await scraper.sample_split(generate_time_s=4.0) + + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(200.0 / 4.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(100.0 / 4.0) + assert not any(k.startswith("vllm/eval/") for k in out) + + @pytest.mark.asyncio async def test_fetch_all_merges_across_nodes_and_aggregates_correctly(): # Two nodes, each running two replicas with disjoint ReplicaId labels. From b82dcfb6894162fd5608c7f10a711ad624c03277 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Tue, 16 Jun 2026 05:35:48 +0000 Subject: [PATCH 16/28] separate eval against train rollout metrics --- skyrl/train/evaluate.py | 10 ++ skyrl/train/trainer.py | 13 +- skyrl/train/utils/vllm_metrics_scraper.py | 138 +++++++++++++++------- tests/train/test_vllm_metrics_scraper.py | 82 +++++++++++++ 4 files changed, 199 insertions(+), 44 deletions(-) diff --git a/skyrl/train/evaluate.py b/skyrl/train/evaluate.py index f24c84f2bd..956075b5a0 100644 --- a/skyrl/train/evaluate.py +++ b/skyrl/train/evaluate.py @@ -1,3 +1,4 @@ +import time from collections import defaultdict from pathlib import Path from typing import Any, Dict, List @@ -58,6 +59,8 @@ async def evaluate( concat_env_extras: List[Dict[str, Any]] = [] concat_uids: List[str] = [] sampling_params = cfg.generator.eval_sampling_params + # Rollout-only time (denominator for the vLLM eval throughput). + eval_generate_time = 0.0 pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress") for _, prompts in enumerate(eval_dataloader): pbar.update(1) @@ -69,7 +72,9 @@ async def evaluate( "eval", global_step, ) + gen_start = time.monotonic() generator_output: GeneratorOutput = await generator.generate(generator_input) + eval_generate_time += time.monotonic() - gen_start validate_generator_output(len(generator_input["prompts"]), generator_output) generator_outputs.append(generator_output) concat_all_envs.extend(generator_input["env_classes"]) @@ -127,6 +132,7 @@ async def evaluate( eval_metrics, ) + eval_metrics["timing/eval_generate"] = eval_generate_time return eval_metrics @@ -160,6 +166,7 @@ async def evaluate_step_wise( concat_env_extras: List[Dict[str, Any]] = [] concat_uids: List[str] = [] sampling_params = cfg.generator.eval_sampling_params + eval_generate_time = 0.0 pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress") for _, prompts in enumerate(eval_dataloader): pbar.update(1) @@ -171,7 +178,9 @@ async def evaluate_step_wise( "eval", global_step, ) + gen_start = time.monotonic() generator_output: GeneratorOutput = await generator.generate(generator_input) + eval_generate_time += time.monotonic() - gen_start traj_id_to_input = { traj_id.instance_id: {"env_class": env_class, "env_extras": env_extra} for traj_id, env_class, env_extra in zip( @@ -244,4 +253,5 @@ async def evaluate_step_wise( eval_metrics, ) + eval_metrics["timing/eval_generate"] = eval_generate_time return eval_metrics diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 851b21b0ba..d2b03d1473 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -426,6 +426,10 @@ async def train(self): or self.global_step == self.total_training_steps ) if force_eval or interval_eval: + # Mark vLLM counters before eval so the eval rollout lands + # in vllm/eval/* rather than blending into the train rollout. + if self._vllm_metrics_scraper is not None: + await self._vllm_metrics_scraper.mark_pre_eval() self._fire("on_eval_start") with Timer("eval", self.all_timings): eval_metrics = await self.eval() @@ -437,11 +441,12 @@ async def train(self): **{f"timing/{k}": v for k, v in self.all_timings.items()}, } if self._vllm_metrics_scraper is not None: - # Engine only generates during the "generate" phase, so - # throughput divides by that, not the full step time. + # vllm/* = train rollout / generate time; on eval steps + # vllm/eval/* = eval rollout / eval rollout time. log_payload.update( - await self._vllm_metrics_scraper.sample( - generation_time_s=self.all_timings.get("generate") + await self._vllm_metrics_scraper.sample_split( + generate_time_s=self.all_timings.get("generate"), + eval_generate_time_s=self.all_metrics.get("timing/eval_generate"), ) ) diff --git a/skyrl/train/utils/vllm_metrics_scraper.py b/skyrl/train/utils/vllm_metrics_scraper.py index 770a912ca4..d63929591c 100644 --- a/skyrl/train/utils/vllm_metrics_scraper.py +++ b/skyrl/train/utils/vllm_metrics_scraper.py @@ -140,9 +140,10 @@ def discover_ray_metrics_urls() -> List[str]: class VLLMMetricsScraper: """Per-step snapshot of selected vLLM metrics from Ray's metrics agents. - The first ``sample()`` call establishes a baseline; rate-style metrics - (throughput, hit rate, average latency) are reported starting from the - second call. + ``sample()`` reports deltas vs. the previous step (used by the fully-async + trainer). ``mark_pre_eval()`` + ``sample_split()`` let the sync trainer + report the train rollout under ``vllm/*`` and the eval rollout under + ``vllm/eval/*`` instead of blending them. """ def __init__( @@ -154,6 +155,8 @@ def __init__( self._timeout = request_timeout_s self._prev_aggregated: Optional[Dict[str, float]] = None self._prev_timestamp: Optional[float] = None + # Counter snapshot taken right before an eval rollout (sync trainer). + self._mid_snapshot: Optional[Dict[str, float]] = None self._client: Optional[httpx.AsyncClient] = None self._warned_empty = False if not self._urls: @@ -194,19 +197,13 @@ async def _fetch_all(self) -> ParsedSamples: merged[key] = value return merged - async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, float]: - """Return a dict of ``vllm/...`` scalars for the current step. - - Empty if no agents are reachable or if no vLLM samples are present - yet (e.g. before any inference has run). + async def _read_snapshot(self) -> Optional[Dict[str, float]]: + """Scrape every agent and reduce to one cumulative value per metric. - ``generation_time_s`` is the time the engine spent generating since the - previous ``sample()`` call; it is the denominator for the throughput - metrics. Pass ``None`` when generation runs for the whole interval - (e.g. fully-async overlap) to fall back to the wall-clock interval. + Returns ``None`` when no endpoints are configured. """ if not self._urls: - return {} + return None parsed = await self._fetch_all() if not parsed and not self._warned_empty: @@ -217,37 +214,95 @@ async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, f ) self._warned_empty = True - now = time.monotonic() - sums = aggregate(parsed, _SUM_METRICS, how="sum") means = aggregate(parsed, _MEAN_METRICS, how="mean") - snapshot = {**sums, **means} + return {**sums, **means} - out: Dict[str, float] = {} - # Instantaneous gauges expose directly. - if _GAUGE_NUM_RUNNING in snapshot: - out["vllm/num_requests_running"] = snapshot[_GAUGE_NUM_RUNNING] - if _GAUGE_NUM_WAITING in snapshot: - out["vllm/num_requests_waiting"] = snapshot[_GAUGE_NUM_WAITING] - if _GAUGE_KV_CACHE_USAGE in snapshot: - out["vllm/kv_cache_usage_perc"] = snapshot[_GAUGE_KV_CACHE_USAGE] - - # Derived metrics need a previous snapshot to take deltas. + async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, float]: + """Return ``vllm/...`` scalars for the current step (empty if unavailable). + + ``generation_time_s`` is the throughput denominator (engine generation + time since the previous call); ``None`` falls back to the wall-clock + interval (fully-async overlap). + """ + snapshot = await self._read_snapshot() + if snapshot is None: + return {} + + now = time.monotonic() if self._prev_aggregated is not None and self._prev_timestamp is not None: dt = max(now - self._prev_timestamp, 1e-9) - # Throughput is per generation-second, not per step-second. - if generation_time_s is not None and generation_time_s > 0: - throughput_window_s = generation_time_s - else: - throughput_window_s = dt - out.update(self._derive(snapshot, self._prev_aggregated, throughput_window_s)) + window = generation_time_s if (generation_time_s is not None and generation_time_s > 0) else dt + out = self._window_metrics(self._prev_aggregated, snapshot, window, "vllm/") + else: + out = self._window_metrics(None, snapshot, None, "vllm/") # gauges only self._prev_aggregated = snapshot self._prev_timestamp = now return out + async def mark_pre_eval(self) -> None: + """Snapshot the counters right before the eval rollout for :meth:`sample_split`.""" + self._mid_snapshot = await self._read_snapshot() + + async def sample_split( + self, + *, + generate_time_s: Optional[float] = None, + eval_generate_time_s: Optional[float] = None, + ) -> Dict[str, float]: + """Sync-trainer sampling that separates train and eval rollouts. + + ``vllm/*`` covers the train rollout (previous step -> pre-eval mark); + when :meth:`mark_pre_eval` was called, ``vllm/eval/*`` covers the eval + rollout (mark -> now). Each is divided by its own generation time. + """ + cur = await self._read_snapshot() + mid = self._mid_snapshot + self._mid_snapshot = None + if cur is None: + return {} + + # Train rollout ends at the pre-eval mark when eval ran, else at cur. + train_end = mid if mid is not None else cur + train_window = generate_time_s if (generate_time_s is not None and generate_time_s > 0) else None + out = self._window_metrics(self._prev_aggregated, train_end, train_window, "vllm/") + + if mid is not None: + eval_window = ( + eval_generate_time_s if (eval_generate_time_s is not None and eval_generate_time_s > 0) else None + ) + out.update(self._window_metrics(mid, cur, eval_window, "vllm/eval/")) + + self._prev_aggregated = cur + return out + + @classmethod + def _window_metrics( + cls, + prev: Optional[Dict[str, float]], + cur: Optional[Dict[str, float]], + throughput_window_s: Optional[float], + prefix: str, + ) -> Dict[str, float]: + """Gauges from ``cur`` plus derived rates over ``cur - prev``.""" + if cur is None: + return {} + out: Dict[str, float] = {} + if _GAUGE_NUM_RUNNING in cur: + out[f"{prefix}num_requests_running"] = cur[_GAUGE_NUM_RUNNING] + if _GAUGE_NUM_WAITING in cur: + out[f"{prefix}num_requests_waiting"] = cur[_GAUGE_NUM_WAITING] + if _GAUGE_KV_CACHE_USAGE in cur: + out[f"{prefix}kv_cache_usage_perc"] = cur[_GAUGE_KV_CACHE_USAGE] + if prev is not None: + out.update(cls._derive(cur, prev, throughput_window_s, prefix)) + return out + @staticmethod - def _derive(cur: Dict[str, float], prev: Dict[str, float], throughput_window_s: float) -> Dict[str, float]: + def _derive( + cur: Dict[str, float], prev: Dict[str, float], throughput_window_s: Optional[float], prefix: str + ) -> Dict[str, float]: out: Dict[str, float] = {} def delta(name: str) -> Optional[float]: @@ -257,27 +312,30 @@ def delta(name: str) -> Optional[float]: # Counter resets (engine restart) shouldn't crash; just skip. return d if d >= 0 else None + # Throughput needs a positive window; the rate/latency metrics don't. + has_window = throughput_window_s is not None and throughput_window_s > 0 + gen_d = delta(_COUNTER_GENERATION_TOKENS) - if gen_d is not None: - out["vllm/generation_throughput_tok_s"] = gen_d / throughput_window_s + if gen_d is not None and has_window: + out[f"{prefix}generation_throughput_tok_s"] = gen_d / throughput_window_s prompt_d = delta(_COUNTER_PROMPT_TOKENS) - if prompt_d is not None: - out["vllm/prompt_throughput_tok_s"] = prompt_d / throughput_window_s + if prompt_d is not None and has_window: + out[f"{prefix}prompt_throughput_tok_s"] = prompt_d / throughput_window_s q_d = delta(_COUNTER_PREFIX_QUERIES) h_d = delta(_COUNTER_PREFIX_HITS) if q_d is not None and h_d is not None and q_d > 0: - out["vllm/prefix_cache_hit_rate"] = h_d / q_d + out[f"{prefix}prefix_cache_hit_rate"] = h_d / q_d ttft_sum_d = delta(_HIST_TTFT_SUM) ttft_count_d = delta(_HIST_TTFT_COUNT) if ttft_sum_d is not None and ttft_count_d is not None and ttft_count_d > 0: - out["vllm/ttft_seconds_avg"] = ttft_sum_d / ttft_count_d + out[f"{prefix}ttft_seconds_avg"] = ttft_sum_d / ttft_count_d itl_sum_d = delta(_HIST_ITL_SUM) itl_count_d = delta(_HIST_ITL_COUNT) if itl_sum_d is not None and itl_count_d is not None and itl_count_d > 0: - out["vllm/tpot_seconds_avg"] = itl_sum_d / itl_count_d + out[f"{prefix}tpot_seconds_avg"] = itl_sum_d / itl_count_d return out diff --git a/tests/train/test_vllm_metrics_scraper.py b/tests/train/test_vllm_metrics_scraper.py index 9066ca4d0f..1307f1e6b1 100644 --- a/tests/train/test_vllm_metrics_scraper.py +++ b/tests/train/test_vllm_metrics_scraper.py @@ -372,6 +372,88 @@ def test_scraper_with_no_urls_is_noop(): assert out == {} +@pytest.mark.asyncio +async def test_sample_split_separates_train_and_eval_rollouts(): + """vllm/* uses the train window; vllm/eval/* uses the eval window. + + baseline gen=500 prompt=1000 + pre-eval gen=600 prompt=1100 (train rollout: +100 gen, +100 prompt) + log-time gen=1100 prompt=1300 (eval rollout: +500 gen, +200 prompt) + with generate_time=2s and eval_generate_time=5s. + """ + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + def snap(gen, prompt): + return _snapshot( + running=1, + waiting=0, + kv=0.4, + prefix_q=0, + prefix_h=0, + prompt_toks=prompt, + gen_toks=gen, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + + texts = iter([snap(500, 1000), snap(600, 1100), snap(1100, 1300)]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + with patch.object(scraper, "_fetch_all", fake_fetch_all): + # Step before eval: establishes the baseline (no throughput yet). + first = await scraper.sample_split(generate_time_s=1.0) + assert "vllm/generation_throughput_tok_s" not in first + + # Eval step: mark before eval, then sample after. + await scraper.mark_pre_eval() + out = await scraper.sample_split(generate_time_s=2.0, eval_generate_time_s=5.0) + + # Train rollout: (600-500)/2 == 50, (1100-1000)/2 == 50. + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(50.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + # Eval rollout: (1100-600)/5 == 100, (1300-1100)/5 == 40. + assert out["vllm/eval/generation_throughput_tok_s"] == pytest.approx(100.0) + assert out["vllm/eval/prompt_throughput_tok_s"] == pytest.approx(40.0) + + +@pytest.mark.asyncio +async def test_sample_split_without_eval_emits_only_train_metrics(): + """On a non-eval step (no mark_pre_eval), no vllm/eval/* keys appear.""" + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + def snap(gen, prompt): + return _snapshot( + running=1, + waiting=0, + kv=0.4, + prefix_q=0, + prefix_h=0, + prompt_toks=prompt, + gen_toks=gen, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + + texts = iter([snap(0, 0), snap(200, 100)]) + + async def fake_fetch_all(): + return parse_metrics_text(next(texts)) + + with patch.object(scraper, "_fetch_all", fake_fetch_all): + await scraper.sample_split(generate_time_s=1.0) + out = await scraper.sample_split(generate_time_s=4.0) + + assert out["vllm/generation_throughput_tok_s"] == pytest.approx(200.0 / 4.0) + assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(100.0 / 4.0) + assert not any(k.startswith("vllm/eval/") for k in out) + + @pytest.mark.asyncio async def test_fetch_all_merges_across_nodes_and_aggregates_correctly(): # Two nodes, each running two replicas with disjoint ReplicaId labels. From b0183f473eda6cbeda9cf9cd880365f89f78fa87 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Tue, 16 Jun 2026 18:02:15 +0000 Subject: [PATCH 17/28] format comments --- skyrl/train/evaluate.py | 1 - skyrl/train/fully_async_trainer.py | 2 -- skyrl/train/trainer.py | 2 -- skyrl/train/utils/vllm_metrics_scraper.py | 1 - 4 files changed, 6 deletions(-) diff --git a/skyrl/train/evaluate.py b/skyrl/train/evaluate.py index 956075b5a0..e0dfc5ad94 100644 --- a/skyrl/train/evaluate.py +++ b/skyrl/train/evaluate.py @@ -59,7 +59,6 @@ async def evaluate( concat_env_extras: List[Dict[str, Any]] = [] concat_uids: List[str] = [] sampling_params = cfg.generator.eval_sampling_params - # Rollout-only time (denominator for the vLLM eval throughput). eval_generate_time = 0.0 pbar = tqdm(total=len(eval_dataloader), initial=0, desc="Evaluation Progress") for _, prompts in enumerate(eval_dataloader): diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index bee8773518..fe31b2fe55 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/skyrl/train/fully_async_trainer.py @@ -459,8 +459,6 @@ async def train(self): timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()} if self._vllm_metrics_scraper is not None: - # Generation overlaps training here, so throughput falls - # back to the wall-clock interval. timing_payload.update(await self._vllm_metrics_scraper.sample()) self.tracker.log(timing_payload, step=self.global_step) self.all_timings = {} diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 24e025def5..4e74a3dbca 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -426,8 +426,6 @@ async def train(self): or self.global_step == self.total_training_steps ) if force_eval or interval_eval: - # Mark vLLM counters before eval so the eval rollout lands - # in vllm/eval/* rather than blending into the train rollout. if self._vllm_metrics_scraper is not None: await self._vllm_metrics_scraper.mark_pre_eval() self._fire("on_eval_start") diff --git a/skyrl/train/utils/vllm_metrics_scraper.py b/skyrl/train/utils/vllm_metrics_scraper.py index d63929591c..bcd4df17ba 100644 --- a/skyrl/train/utils/vllm_metrics_scraper.py +++ b/skyrl/train/utils/vllm_metrics_scraper.py @@ -155,7 +155,6 @@ def __init__( self._timeout = request_timeout_s self._prev_aggregated: Optional[Dict[str, float]] = None self._prev_timestamp: Optional[float] = None - # Counter snapshot taken right before an eval rollout (sync trainer). self._mid_snapshot: Optional[Dict[str, float]] = None self._client: Optional[httpx.AsyncClient] = None self._warned_empty = False From 082e4714bd49be937a7da91517cfb9809b514f13 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Tue, 16 Jun 2026 21:59:13 +0000 Subject: [PATCH 18/28] remove legacy code --- ...tron_dapo_qwen3.5_9b_specdecode_notrain.sh | 2 +- skyrl/backends/skyrl_train/mtp/__init__.py | 36 +----- skyrl/backends/skyrl_train/mtp/adapter.py | 33 +---- .../skyrl_train/mtp/draft_loss_wrapper.py | 122 ------------------ .../backends/skyrl_train/utils/torch_utils.py | 15 --- .../workers/megatron/megatron_worker.py | 4 +- skyrl/train/config/config.py | 14 +- .../skyrl_train/mtp/test_mtp_config.py | 8 -- .../skyrl_train/mtp/test_mtp_torch_utils.py | 16 --- .../backends/skyrl_train/mtp/test_soft_ce.py | 66 +--------- 10 files changed, 20 insertions(+), 296 deletions(-) delete mode 100644 skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh index b5a6731ab4..9b7b0975a7 100644 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh +++ b/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh @@ -203,7 +203,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo" \ + trainer.project_name="qwen3_5_dapo_2" \ trainer.run_name="sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.export_path="/mnt/local_storage/exports/sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ trainer.hf_save_interval=300 \ diff --git a/skyrl/backends/skyrl_train/mtp/__init__.py b/skyrl/backends/skyrl_train/mtp/__init__.py index 0949ee010f..08b5f3754b 100644 --- a/skyrl/backends/skyrl_train/mtp/__init__.py +++ b/skyrl/backends/skyrl_train/mtp/__init__.py @@ -1,32 +1,10 @@ # Decoupled Multi-Token Prediction (MTP) draft-head training. # -# This package mirrors the design used by NeMo-RL's draft-model training -# (https://github.com/NVIDIA-NeMo/RL, ``nemo_rl/models/megatron/draft`` and -# ``nemo_rl/algorithms/loss``): the trunk's hidden states are *detached* before -# the draft/MTP head runs (decoupling the draft gradient from the policy -# backbone), the head is supervised by an explicit, configurable loss -# (soft cross-entropy distillation against the policy's own next-token -# distribution, or hard next-token cross-entropy), and the two losses are -# combined legibly as ``policy_loss + w * draft_loss`` instead of being -# entangled inside Megatron's ``MTPLossAutoScaler``. +# The trunk's hidden states are *detached* before the draft/MTP head runs +# (decoupling the draft gradient from the policy backbone), and the head is +# supervised by an explicit, configurable loss (soft cross-entropy distillation +# against the policy's own next-token distribution, or hard next-token +# cross-entropy) instead of being entangled inside Megatron's MTPLossAutoScaler. # -# The pieces here are backend-agnostic. Only the ``adapter`` wiring is -# backend-specific (Megatron today; FSDP is a planned follow-up). - -from skyrl.backends.skyrl_train.mtp.draft_loss_wrapper import ( - DraftLossWrapper, - combine_policy_and_draft_loss, -) -from skyrl.backends.skyrl_train.mtp.soft_ce import ( - build_teacher_logits, - draft_hard_ce, - draft_soft_ce, -) - -__all__ = [ - "DraftLossWrapper", - "combine_policy_and_draft_loss", - "build_teacher_logits", - "draft_hard_ce", - "draft_soft_ce", -] +# Consumers import directly from the submodules (``soft_ce``, ``hidden_capture``, +# ``adapter``); this package has no re-exports of its own. diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py index db8750230a..598d365484 100644 --- a/skyrl/backends/skyrl_train/mtp/adapter.py +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -1,24 +1,14 @@ -# Backend seam for decoupled MTP/draft-head training. +# Output-head projection for decoupled MTP/draft-head training. # -# The loss (``soft_ce``), the combination (``draft_loss_wrapper``) and the -# capture mechanism (``hidden_capture``) are backend-agnostic. Only two things -# differ per backend / draft style: (a) where the trunk hidden states are -# captured and (b) how the draft head turns those hidden states into vocab -# logits. ``DraftAdapter`` pins down that contract. -# -# Implemented today: native-MTP draft training on Megatron (capture the model's -# built-in ``self.mtp`` block and project through the shared output layer), which +# The loss (``soft_ce``) and the capture mechanism (``hidden_capture``) are +# backend-agnostic. This module holds the Megatron-specific piece: turning the +# captured MTP hidden states into vocab logits via the shared output layer. It # works for any model that ships native MTP heads regardless of base class # (``GPTModel`` for DeepSeek/GLM/Qwen3-Next, ``MambaModel`` for Qwen3.5/NemotronH). -# -# Planned follow-ups implement the SAME protocol and reuse the loss/combination -# unchanged: (1) a NeMo-RL-style Eagle adapter that captures auxiliary policy -# hidden states and projects them through a separate Eagle draft model (for models -# WITHOUT native MTP heads); (2) an FSDP backend adapter. from __future__ import annotations -from typing import List, Protocol, runtime_checkable +from typing import List import torch @@ -51,19 +41,6 @@ def backward(ctx, grad): return grad.contiguous() -@runtime_checkable -class DraftAdapter(Protocol): - """Per-backend / per-draft-style hooks for decoupled MTP/draft training. - - An Eagle adapter would implement the same two methods: ``capture_context`` - grabs the auxiliary policy hidden states, and ``project_to_logits`` runs the - (separate) draft model's output head.""" - - def capture_context(self, model, detach_trunk: bool): ... - - def project_to_logits(self, hidden_states_per_layer, model) -> List: ... - - def project_mtp_hidden_to_logits(hidden_states_per_layer, model, detach_output_weight: bool = False) -> List: """Run the model's shared output layer on each captured MTP hidden-state chunk. diff --git a/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py b/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py deleted file mode 100644 index 9e519914f1..0000000000 --- a/skyrl/backends/skyrl_train/mtp/draft_loss_wrapper.py +++ /dev/null @@ -1,122 +0,0 @@ -# Combine the policy loss with the decoupled draft (MTP) loss. -# -# Mirrors NeMo-RL's ``DraftLossWrapper`` -# (https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/wrapper.py): -# the policy loss and the draft loss are computed independently and combined as -# ``policy_loss + w * draft_loss``. Because the draft head is fed *detached* -# trunk hidden states and a *detached* teacher distribution, the draft term only -# produces gradients for the draft-head parameters — the addition is just a -# convenient single backward, not an entanglement of the two objectives. - -from dataclasses import dataclass -from typing import Optional - -import torch - -from skyrl.backends.skyrl_train.mtp.soft_ce import ( - build_teacher_logits, - draft_hard_ce, - draft_soft_ce, - shift_mask_for_mtp, -) - - -@dataclass -class DraftLossConfig: - """Resolved knobs for the draft loss (populated from MegatronConfig).""" - - loss_weight: float = 0.1 - loss_type: str = "soft_ce" # "soft_ce" | "hard_ce" - - -def combine_policy_and_draft_loss( - policy_loss: torch.Tensor, - student_logits_per_layer: list[torch.Tensor], - main_logits: torch.Tensor, - mask: torch.Tensor, - cfg: DraftLossConfig, - hard_labels: Optional[torch.Tensor] = None, - global_normalization_factor: Optional[torch.Tensor] = None, - vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None, - vocab_start_index: Optional[int] = None, -) -> tuple[torch.Tensor, dict]: - """Return ``policy_loss + w * draft_loss`` and a metrics dict. - - Args: - policy_loss: The already-computed scalar policy loss. - student_logits_per_layer: One ``[batch, seq, vocab(/tp)]`` draft-logits - tensor per MTP depth, aligned (same de-padding) with ``main_logits``. - main_logits: ``[batch, seq, vocab(/tp)]`` policy logits — the soft-CE - teacher source. - mask: ``[batch, seq]`` token mask over the region the draft is trained on. - cfg: Draft loss configuration. - hard_labels: ``[batch, seq]`` base next-token labels (``seq[t+1]`` at ``t``). - Rolled internally per MTP depth. Required when ``cfg.loss_type == "hard_ce"``. - global_normalization_factor: Optional global valid-token count for the - masked-mean denominator (cross-microbatch / DP correct reduction). - vocab_parallel_group: TP group when logits are vocab-sharded. - vocab_start_index: This rank's vocab offset (required for vocab-parallel - hard CE). - - Returns: - ``(combined_loss, metrics)`` where ``metrics`` carries the (detached) - draft loss for logging. - """ - draft_losses = [] - for layer_idx, student_logits in enumerate(student_logits_per_layer): - # Hard CE's target is the *token* ``seq[t+k+2]`` — one position past the soft-CE teacher - # (a *distribution* read at ``t+k+1``) — so its validity mask needs one extra shift, or the - # last in-range position of every row trains against a zeroed/pad label. - mask_shift = layer_idx + 1 if cfg.loss_type == "hard_ce" else layer_idx - layer_mask = shift_mask_for_mtp(mask, mask_shift) - if cfg.loss_type == "hard_ce": - assert hard_labels is not None, "hard_ce requires hard_labels" - layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) - draft_losses.append( - draft_hard_ce( - student_logits, - layer_labels, - layer_mask, - global_normalization_factor=global_normalization_factor, - vocab_parallel_group=vocab_parallel_group, - vocab_start_index=vocab_start_index, - ) - ) - else: - teacher_logits = build_teacher_logits(main_logits, layer_idx, detach=True) - draft_losses.append( - draft_soft_ce( - student_logits, - teacher_logits, - layer_mask, - global_normalization_factor=global_normalization_factor, - vocab_parallel_group=vocab_parallel_group, - ) - ) - - draft_loss = torch.stack(draft_losses).mean() - combined = policy_loss + cfg.loss_weight * draft_loss - metrics = { - "draft_loss": draft_loss.detach().item(), - "mtp_loss": draft_loss.detach().item(), # alias kept for existing dashboards - } - return combined, metrics - - -class DraftLossWrapper: - """Object form of :func:`combine_policy_and_draft_loss` for backends that - prefer a stateful seam (matches NeMo-RL's ``DraftLossWrapper`` shape). The - FSDP backend will reuse this unchanged once its adapter lands.""" - - def __init__(self, cfg: DraftLossConfig): - self.cfg = cfg - - def __call__(self, policy_loss, student_logits_per_layer, main_logits, mask, **kwargs): - return combine_policy_and_draft_loss( - policy_loss, - student_logits_per_layer, - main_logits, - mask, - self.cfg, - **kwargs, - ) diff --git a/skyrl/backends/skyrl_train/utils/torch_utils.py b/skyrl/backends/skyrl_train/utils/torch_utils.py index 2ac21291c2..d7590298cb 100644 --- a/skyrl/backends/skyrl_train/utils/torch_utils.py +++ b/skyrl/backends/skyrl_train/utils/torch_utils.py @@ -207,21 +207,6 @@ def build_mtp_next_token_labels(sequences: torch.Tensor) -> torch.Tensor: return labels -def build_mtp_loss_mask(attention_mask: torch.Tensor) -> torch.Tensor: - """Build the MTP loss mask from an attention mask. - - We train the MTP heads on every real token of the sequence (next-token prediction over the - full prompt+response), mirroring standard MTP pretraining. ``attention_mask`` already marks - real (non-pad) tokens, so it doubles as the loss mask. Per-sequence boundaries (the final - real token, whose MTP target is invalid) are handled by Megatron's roll in - ``process_mtp_loss`` and need no special treatment here. - - Returns a float tensor so it can flow through the same de-padding/packing transforms applied - to the token ids. - """ - return attention_mask.to(torch.float32) - - def safe_exp_delta(delta: torch.Tensor, clip: float = 20.0, out_dtype=None) -> torch.Tensor: """ Clamp the delta before exponentiating to avoid potential overflow. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 0b1088f48e..ba494110d8 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -442,8 +442,8 @@ def init_configs( # The native MTP heads are still *built* (so the megatron-bridge weight mappings # round-trip them to HF / vLLM), but SkyRL trains them with a decoupled, explicit # loss (see MegatronModelWrapper) instead of Megatron's in-forward - # process_mtp_loss / MTPLossAutoScaler path. We therefore do not set - # provider.mtp_loss_scaling_factor; the explicit weight is megatron_config.mtp_loss_weight. + # process_mtp_loss / MTPLossAutoScaler path. The explicit weight is + # megatron_config.mtp_loss_weight. logger.info( f"MTP enabled (decoupled): mtp_num_layers={provider.mtp_num_layers}, " f"mtp_loss_type={megatron_config.mtp_loss_type}, " diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 366fe6cda0..2e4bb2a34c 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -208,10 +208,10 @@ class MegatronConfig(BaseConfig): mtp_loss_weight: float = 0.1 """Weight ``w`` of the decoupled MTP/draft loss in ``policy_loss + w * draft_loss``. - Replaces the role of Megatron's opaque ``mtp_loss_scaling_factor`` backward-scale. SkyRL now - trains the native MTP heads with an *explicit* loss on *detached* trunk hidden states (so the - draft gradient never pulls on the policy backbone) and adds it to the policy loss with this - weight. Only used when MTP layers are active.""" + SkyRL trains the native MTP heads with an *explicit* loss on *detached* trunk hidden states (so + the draft gradient never pulls on the policy backbone) and adds it to the policy loss with this + weight, instead of Megatron's opaque in-forward MTP backward-scale. Only used when MTP layers + are active.""" mtp_loss_type: str = "soft_ce" """Supervision for the MTP/draft head: @@ -231,9 +231,6 @@ class MegatronConfig(BaseConfig): head means the draft loss directly nudges the policy's own logits every step. Matches slime's MTP-RL training (only ``.mtp.`` params receive gradients). Set False for NeMo-RL's behaviour (shared embedding/head also trained by the draft loss).""" - mtp_loss_scaling_factor: Optional[float] = None - """Deprecated alias for ``mtp_loss_weight`` (kept for back-compat). If set, it overrides - ``mtp_loss_weight``.""" mtp_loss_chunk_size: Optional[int] = 1024 """Sequence-chunk size for the decoupled MTP/draft loss. The draft loss materializes full ``[batch, seq, vocab]`` softmax tensors; at large vocab (e.g. Qwen3.5's 248K) computing the whole @@ -251,9 +248,6 @@ class MegatronConfig(BaseConfig): Typical: 64-128 (acceptance is governed by the top tokens, so the dropped tail is benign).""" def __post_init__(self): - # Back-compat: the old `mtp_loss_scaling_factor` knob now maps onto the explicit loss weight. - if self.mtp_loss_scaling_factor is not None: - self.mtp_loss_weight = self.mtp_loss_scaling_factor # Backfill defaults for any keys the user didn't override so an override dict # doesn't have to repeat every default just to set one value. if self.transformer_config_kwargs is None: diff --git a/tests/backends/skyrl_train/mtp/test_mtp_config.py b/tests/backends/skyrl_train/mtp/test_mtp_config.py index 32b4ca03b5..579e0a4201 100644 --- a/tests/backends/skyrl_train/mtp/test_mtp_config.py +++ b/tests/backends/skyrl_train/mtp/test_mtp_config.py @@ -27,8 +27,6 @@ def test_megatron_config_mtp_defaults(): # embedding/output weight (== the lm_head on tied-embedding models like Qwen3.5) is detached # in both the output projection and the MTP block's re-embedding (slime-style). assert cfg.mtp_detach_shared_output is True - # Deprecated alias is unset by default. - assert cfg.mtp_loss_scaling_factor is None def test_megatron_config_mtp_overrides_parse(): @@ -40,12 +38,6 @@ def test_megatron_config_mtp_overrides_parse(): assert cfg.mtp_loss_type == "hard_ce" -def test_megatron_config_mtp_loss_scaling_factor_back_compat(): - # The deprecated scaling-factor knob maps onto the explicit loss weight. - cfg = build_nested_dataclass(MegatronConfig, {"mtp_loss_scaling_factor": 0.25}) - assert cfg.mtp_loss_weight == 0.25 - - def test_megatron_config_mtp_force_disable(): # An explicit 0 is how a user force-disables MTP even on an MTP-capable model. cfg = build_nested_dataclass(MegatronConfig, {"mtp_num_layers": 0}) diff --git a/tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py b/tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py index dee9a9068d..d7f2c74f3c 100644 --- a/tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py +++ b/tests/backends/skyrl_train/mtp/test_mtp_torch_utils.py @@ -6,7 +6,6 @@ import torch from skyrl.backends.skyrl_train.utils.torch_utils import ( - build_mtp_loss_mask, build_mtp_next_token_labels, ) @@ -35,18 +34,3 @@ def test_build_mtp_next_token_labels_no_wraparound_into_first_token(): labels = build_mtp_next_token_labels(sequences) assert labels[0, -1].item() == 0 assert labels[0, -1].item() != sequences[0, 0].item() or sequences[0, 0].item() == 0 - - -def test_build_mtp_loss_mask_from_attention_mask(): - attention_mask = torch.tensor([[1, 1, 1, 0], [1, 1, 0, 0]]) - loss_mask = build_mtp_loss_mask(attention_mask) - assert loss_mask.dtype == torch.float32 - assert torch.equal(loss_mask, attention_mask.to(torch.float32)) - - -def test_build_mtp_loss_mask_accepts_bool_mask(): - # forward_step passes attention_mask.to(bool); the helper must still produce a float mask. - attention_mask = torch.tensor([[True, True, False]]) - loss_mask = build_mtp_loss_mask(attention_mask) - assert loss_mask.dtype == torch.float32 - assert torch.equal(loss_mask, torch.tensor([[1.0, 1.0, 0.0]])) diff --git a/tests/backends/skyrl_train/mtp/test_soft_ce.py b/tests/backends/skyrl_train/mtp/test_soft_ce.py index 0437d5b59a..e5fd36c2db 100644 --- a/tests/backends/skyrl_train/mtp/test_soft_ce.py +++ b/tests/backends/skyrl_train/mtp/test_soft_ce.py @@ -1,4 +1,4 @@ -"""CPU unit tests for the decoupled MTP draft losses and loss combination. +"""CPU unit tests for the decoupled MTP draft losses. uv run --isolated --extra dev pytest tests/backends/skyrl_train/mtp/test_soft_ce.py """ @@ -6,10 +6,6 @@ import torch import torch.nn.functional as F -from skyrl.backends.skyrl_train.mtp.draft_loss_wrapper import ( - DraftLossConfig, - combine_policy_and_draft_loss, -) from skyrl.backends.skyrl_train.mtp.soft_ce import ( build_teacher_logits, draft_hard_ce, @@ -297,66 +293,6 @@ def test_left_pad_zero_logits_do_not_inflate_loss(): assert loss.item() < torch.log(torch.tensor(float(V))).item() -def test_combine_is_policy_plus_weighted_draft(): - torch.manual_seed(3) - policy_loss = torch.tensor(2.0, requires_grad=True) - student = torch.randn(2, 5, 7, requires_grad=True) - main_logits = torch.randn(2, 5, 7) - mask = torch.ones(2, 5) - cfg = DraftLossConfig(loss_weight=0.5, loss_type="soft_ce") - - combined, metrics = combine_policy_and_draft_loss(policy_loss, [student], main_logits, mask, cfg) - # Recompute the draft term independently and check the combination identity. - teacher = build_teacher_logits(main_logits, 0) - draft = draft_soft_ce(student, teacher, shift_mask_for_mtp(mask, 0)) - assert torch.allclose(combined, policy_loss + 0.5 * draft, atol=1e-6) - assert abs(metrics["mtp_loss"] - draft.item()) < 1e-6 - - -def test_combine_hard_ce_uses_rolled_labels(): - torch.manual_seed(4) - policy_loss = torch.tensor(1.0) - student = torch.randn(1, 5, 7, requires_grad=True) - main_logits = torch.randn(1, 5, 7) - mask = torch.ones(1, 5) - base_labels = torch.randint(0, 7, (1, 5)) - cfg = DraftLossConfig(loss_weight=1.0, loss_type="hard_ce") - - combined, _ = combine_policy_and_draft_loss(policy_loss, [student], main_logits, mask, cfg, hard_labels=base_labels) - layer_labels = torch.roll(base_labels, shifts=-1, dims=1) - # Hard CE masks with one extra shift (label at t+k+2 must be a real token), unlike soft CE. - draft = draft_hard_ce(student, layer_labels, shift_mask_for_mtp(mask, 1)) - assert torch.allclose(combined, policy_loss + draft, atol=1e-6) - - -def test_combine_hard_ce_excludes_boundary_with_invalid_label(): - # Regression: hard CE at depth k supervises position t against the TOKEN seq[t+k+2] — one - # position further than the soft-CE teacher (a distribution read at t+k+1). With the soft-CE - # mask shift, the last in-range position (t = last_real - (k+1)) was trained against the - # zeroed boundary label (token id 0). Build a student that is perfectly confident in the true - # label at every genuinely-valid position but puts ~zero mass on token 0 at the boundary - # position: the loss must be ~0 (boundary excluded), not inflated by -log q(0). - S, V = 6, 7 - seq = torch.arange(1, S + 1) % V # tokens 1..6 — never 0, so label 0 only comes from the boundary - hard_labels = torch.roll(seq, -1).unsqueeze(0).clone() # labels[t] = seq[t+1] - hard_labels[:, -1] = 0 - mask = torch.ones(1, S) - - # depth 0: true label at t is seq[t+2]; valid sources are t = 0..S-3. - student = torch.full((1, S, V), -30.0) - for t in range(S - 2): - student[0, t, seq[t + 2]] = 30.0 - # boundary position t = S-2: confident in a NON-zero token, so a leaked label-0 target - # would contribute ~60 nats. - student[0, S - 2, 1] = 30.0 - - cfg = DraftLossConfig(loss_weight=1.0, loss_type="hard_ce") - combined, metrics = combine_policy_and_draft_loss( - torch.tensor(0.0), [student], torch.zeros(1, S, V), mask, cfg, hard_labels=hard_labels - ) - assert metrics["draft_loss"] < 1e-3, metrics["draft_loss"] - - def test_unpadded_vocab_shard_width(): from skyrl.backends.skyrl_train.mtp.soft_ce import unpadded_vocab_shard_width From 5802b9083180f73ff3f3378ac146ee5e908268f8 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 17 Jun 2026 00:06:37 +0000 Subject: [PATCH 19/28] reduce comments --- skyrl/backends/skyrl_train/mtp/__init__.py | 12 +- skyrl/backends/skyrl_train/mtp/adapter.py | 46 ++---- .../skyrl_train/mtp/hidden_capture.py | 86 ++++------- skyrl/backends/skyrl_train/mtp/soft_ce.py | 138 ++++++------------ .../megatron/megatron_model_wrapper.py | 69 +++------ skyrl/train/config/config.py | 100 +++++-------- 6 files changed, 145 insertions(+), 306 deletions(-) diff --git a/skyrl/backends/skyrl_train/mtp/__init__.py b/skyrl/backends/skyrl_train/mtp/__init__.py index 08b5f3754b..e71d2c31c5 100644 --- a/skyrl/backends/skyrl_train/mtp/__init__.py +++ b/skyrl/backends/skyrl_train/mtp/__init__.py @@ -1,10 +1,6 @@ # Decoupled Multi-Token Prediction (MTP) draft-head training. # -# The trunk's hidden states are *detached* before the draft/MTP head runs -# (decoupling the draft gradient from the policy backbone), and the head is -# supervised by an explicit, configurable loss (soft cross-entropy distillation -# against the policy's own next-token distribution, or hard next-token -# cross-entropy) instead of being entangled inside Megatron's MTPLossAutoScaler. -# -# Consumers import directly from the submodules (``soft_ce``, ``hidden_capture``, -# ``adapter``); this package has no re-exports of its own. +# The trunk hidden states are detached before the MTP head runs, and the head is supervised by an +# explicit loss (soft-CE distillation against the policy's own next-token distribution, or hard CE) +# rather than Megatron's MTPLossAutoScaler. Consumers import directly from the submodules +# (``soft_ce``, ``hidden_capture``, ``adapter``); this package has no re-exports of its own. diff --git a/skyrl/backends/skyrl_train/mtp/adapter.py b/skyrl/backends/skyrl_train/mtp/adapter.py index 598d365484..b0373d9cbe 100644 --- a/skyrl/backends/skyrl_train/mtp/adapter.py +++ b/skyrl/backends/skyrl_train/mtp/adapter.py @@ -16,18 +16,11 @@ class _CanonicalGradStrides(torch.autograd.Function): """Identity forward; backward hands upstream a stride-canonical gradient. - The draft projection runs through megatron's ``LinearWithFrozenWeight`` (the - detached-weight path), whose hand-written dgrad is ``grad_output.matmul(weight)``. - ``torch.matmul`` only folds that 3D x 2D product into a flat ``mm`` when the grad - passes ``should_fold``'s stride check, which does NOT skip size-1 dims. With - micro-batch 1, the grad reaching it is a transpose-backward view whose size-1 batch - dim carries a stale stride (the adapter's ``.contiguous()`` no-ops on such views), - so matmul falls back to a broadcast ``bmm`` (batch=seq, M=1) that runs ~100x slower - than the equivalent ``mm`` — measured 1.27s vs 10ms per microbatch at - [8344, 1, 62080] x [62080, 4096] on H100. Re-viewing the (dense) grad rewrites the - strides to canonical form at zero copy and restores the ``mm`` dispatch. The - non-detached path is unaffected either way (``should_fold`` folds early when the - weight requires grad), and the no-copy view keeps this shim free there too. + The detached-weight projection (``LinearWithFrozenWeight``) backprops via ``grad.matmul(weight)``. + At micro-batch 1 the grad is a transpose-backward view whose stale size-1 batch stride fails + matmul's ``should_fold`` check, so it dispatches a broadcast ``bmm`` ~100x slower than the + equivalent ``mm`` (measured 1.27s vs 10ms/microbatch on H100). Re-viewing the dense grad restores + canonical strides at zero copy. No-op on the non-detached path. """ @staticmethod @@ -44,32 +37,21 @@ def backward(ctx, grad): def project_mtp_hidden_to_logits(hidden_states_per_layer, model, detach_output_weight: bool = False) -> List: """Run the model's shared output layer on each captured MTP hidden-state chunk. - Mirrors the model's own ``_postprocess``: logits come out of the output layer - in ``[seq, batch, vocab/tp]`` and are transposed to ``[batch, seq, vocab/tp]`` - to match the main-logits layout. The returned tensors are still in Megatron's - *internal* (packed / left-removed) sequence layout; the caller applies the - same de-padding transform used for the main logits so they align with the - teacher. + Like the model's ``_postprocess``: output-layer logits come out ``[seq, batch, vocab/tp]`` and are + transposed to ``[batch, seq, vocab/tp]``. The result is still in Megatron's internal (packed / + left-removed) layout; the caller de-pads it with the same transform as the main logits. - Args: - hidden_states_per_layer: list of ``[seq, batch, hidden]`` tensors, one per - MTP depth (from ``MTPHiddenCapture.compute_student_hidden_states``). - model: the unwrapped Megatron model (``GPTModel``/``MambaModel``/...), - which exposes ``output_layer`` and, when embeddings are tied, - ``shared_embedding_or_output_weight``. - - Returns: - list of ``[batch, seq, vocab/tp]`` student-logits tensors (internal layout). + Takes one ``[seq, batch, hidden]`` tensor per MTP depth (from + ``MTPHiddenCapture.compute_student_hidden_states``) and returns one ``[batch, seq, vocab/tp]`` + student-logits tensor each. """ output_weight = None if getattr(model, "share_embeddings_and_output_weights", False): output_weight = model.shared_embedding_or_output_weight() if detach_output_weight: - # Isolate the output projection from the draft gradient: detach the shared/tied weight, - # or — for untied models — pass the output layer's own weight explicitly as a detached - # tensor (ColumnParallelLinear uses a provided ``weight`` verbatim). Mirrors slime's - # MTP-RL megatron patch. Combined with the capture's detached re-embedding, the draft - # loss then trains only the MTP-head parameters. + # Isolate the output projection from the draft gradient: detach the shared/tied weight, or for + # untied models pass the output layer's own weight as a detached tensor. With the capture's + # detached re-embedding, the draft loss then trains only the MTP-head parameters. if output_weight is None: output_weight = getattr(model.output_layer, "weight", None) if output_weight is not None: diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index 60538ae229..633f019569 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -1,35 +1,16 @@ -# Decoupled capture of the trunk hidden states that feed the MTP/draft head. +# Decoupled capture of the trunk hidden states that feed the MTP/draft head. Inspired by NeMo-RL's +# draft-model training (https://github.com/NVIDIA-NeMo/RL), adapted for SkyRL's Megatron backend. # -# This is the SkyRL analogue of NeMo-RL's ``HiddenStateCapture`` -# (https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/models/megatron/draft/hidden_capture.py). -# -# Works for any Megatron model that builds native MTP heads and exposes ``self.mtp`` (the shared -# ``MultiTokenPredictionBlock``): ``GPTModel`` (DeepSeek-V3 / GLM / Qwen3-Next, ...) and -# ``MambaModel`` (Qwen3.5 / NemotronH, ...). Both run ``self.mtp(...)`` + ``process_mtp_loss(...)`` -# inside ``forward`` and surface the same output layout, so the capture below is model-agnostic. -# -# Megatron runs its native MTP block (``self.mtp``) inside the forward whenever ``self.mtp_process`` -# is set. Crucially, ``MultiTokenPredictionBlock.forward`` passes the trunk hidden states through as -# a passthrough chunk of its output, and ``process_mtp_loss`` returns that chunk as the hidden states -# for the *main* logits. So we must NOT tamper with the in-forward call: the main policy logits have -# to stay connected to the trunk for the policy loss to train it. -# -# Instead we: -# 1. register a forward pre-hook on ``self.mtp`` that simply *records* the keyword arguments -# Megatron built for it (input_ids, position_ids, the live trunk hidden states, rotary -# embeddings, attention mask, packed-seq params, the shared embedding, ...). The pre-hook does -# not modify the call, so the in-forward MTP run proceeds normally and the main logits stay -# coupled to the trunk. -# 2. after the forward, re-invoke ``self.mtp`` with the recorded arguments but with the trunk -# hidden states *detached* (when ``detach_trunk``). This second pass is the decoupled draft -# forward: its gradient reaches the MTP-head (and shared embedding/output) parameters but never -# the policy backbone. Re-using the recorded kwargs means we never have to reconstruct rotary -# embeddings or attention masks ourselves. -# -# The in-forward MTP layers are recomputed (their output is discarded by ``process_mtp_loss`` when -# no labels are passed), so there is a small extra-compute cost — one MTP block forward, typically a -# single transformer layer. This is a deliberate trade for robustness over reaching into GPTModel's -# rotary/mask construction. +# Works for any Megatron model with native MTP heads exposing ``self.mtp`` (the shared +# ``MultiTokenPredictionBlock``): ``GPTModel`` (DeepSeek-V3 / GLM / Qwen3-Next) and ``MambaModel`` +# (Qwen3.5 / NemotronH). The in-forward ``self.mtp(...)`` passes the trunk hidden states through to +# ``process_mtp_loss`` as the main logits, so we must not tamper with it. Instead: +# 1. a forward pre-hook records the kwargs Megatron built for ``self.mtp`` (without modifying them); +# 2. after the forward, we re-invoke ``self.mtp`` with those kwargs but the trunk hidden states +# detached -- the decoupled draft pass, whose gradient reaches the MTP head but never the policy +# backbone. Reusing the kwargs avoids reconstructing rotary embeddings / attention masks. +# The replay costs one extra MTP-block forward (typically a single layer) -- a deliberate trade for +# robustness over reaching into GPTModel's rotary/mask construction. from __future__ import annotations @@ -45,14 +26,11 @@ def _unwrap_model(model): def _resolve_mtp_host(model): - """Return the submodule that actually owns the native MTP block (``.mtp``). - - For a plain ``GPTModel`` / ``MambaModel`` that is the model itself. For vision-language wrappers - (e.g. Qwen3.5-VL's ``Qwen3VLModel``) the text backbone *and* its MTP head are nested one level - down at ``model.language_model`` (the bridge maps the heads to ``language_model.mtp.*``), so the - top-level model has no ``.mtp`` and the capture would otherwise silently find nothing. We descend - through the common ``.language_model`` nesting until we find a module exposing ``.mtp``; if none - does, return the top-level model unchanged (capture then yields ``None``, as before).""" + """Return the submodule that owns the native MTP block (``.mtp``). + + Usually the model itself, but for VL wrappers (e.g. Qwen3.5-VL) the MTP head is nested at + ``model.language_model``. Descend the ``.language_model`` chain until we find a module with ``.mtp``; + if none, return the model unchanged (capture then yields ``None``).""" seen = set() cur = model for _ in range(4): # bounded descent guard against pathological nesting / cycles @@ -93,18 +71,13 @@ class MTPHiddenCapture: """ def __init__(self, model, detach_trunk: bool = True, detach_shared_embedding: bool = False): - # Resolve the module that owns the MTP block: the unwrapped model itself for plain - # GPT/Mamba models, or the nested ``.language_model`` for VL wrappers (e.g. Qwen3.5-VL). - # capture.model is also what the caller projects through (output_layer / shared weight), - # so it must be the same module that holds the heads. + # The module that owns the MTP block is also what the caller projects through (output_layer / + # shared weight), so they must match. self.model = _resolve_mtp_host(_unwrap_model(model)) self.detach_trunk = detach_trunk - # The MTP block re-embeds the rolled input ids through the model's shared embedding - # (``embedding=self.embedding`` in its kwargs) — a second gradient path into the shared - # embedding weight, separate from the output projection. When the caller wants the shared - # weights fully isolated from the draft loss (``mtp_detach_shared_output``), the replay must - # detach this path too, else tied-embedding models still train the lm_head through the - # lookup (slime's MTP-RL patch detaches the same two paths). + # The MTP block re-embeds the rolled input ids through the shared embedding -- a second + # gradient path into the embedding weight besides the output projection. detach_shared_embedding + # cuts it too, so tied-embedding models don't keep training the lm_head through the lookup. self.detach_shared_embedding = detach_shared_embedding self._args = None self._kwargs = None @@ -130,11 +103,9 @@ def capture(self): return self._args = None self._kwargs = None - # Run the MTP block in eval mode for both the in-forward pass and the replay. Megatron's - # full-activation-recompute path (recompute_granularity='full' and module.training) routes - # through CheckpointFunction, which cannot save the non-tensor PackedSeqParams for backward. - # Eval skips that path (dropout is 0 in these configs; gradients still flow), and MTP is a - # single tiny layer so keeping its activations is negligible. + # Run the MTP block in eval mode (forward + replay): Megatron's full-recompute path routes + # through CheckpointFunction, which can't save the non-tensor PackedSeqParams for backward. + # Eval skips it (dropout is 0 here; gradients still flow) and MTP is one tiny layer. self._prev_training = mtp.training mtp.eval() self._handles.append(mtp.register_forward_pre_hook(self._pre_hook, with_kwargs=True)) @@ -159,10 +130,9 @@ def compute_student_hidden_states(self) -> Optional[List]: kwargs = dict(self._kwargs) if self.detach_trunk: hidden = kwargs.get("hidden_states") - # The detach below only patches the *keyword* argument. Megatron currently calls - # ``self.mtp(...)`` with kwargs only (GPTModel + MambaModel); if a future version passes - # hidden_states positionally, the detach would silently no-op and the draft gradient - # would couple back into the policy trunk — fail loudly instead. + # We only patch the keyword arg. Megatron passes hidden_states as a kwarg today; if a + # future version passes it positionally the detach would silently no-op (re-coupling the + # trunk), so fail loudly instead. assert hidden is not None, ( "MTP capture: 'hidden_states' was not passed to the MTP block as a keyword argument " "(megatron-core call convention changed?); cannot detach the trunk for decoupled " diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index 376e0fbb2b..bceab99693 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -1,14 +1,9 @@ # Explicit draft-head losses for decoupled Multi-Token Prediction (MTP). # -# Ported from NeMo-RL's ``DraftCrossEntropyLossFn`` and the ``DRAFT`` branch of -# ``prepare_loss_input``: -# https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/loss_functions.py -# https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/algorithms/loss/utils.py -# -# The soft cross-entropy matches the forward-KL student gradient -# (``softmax(student) - softmax(teacher)``) and is the default supervision for -# the draft head: the teacher is the *policy's own* next-token distribution -# (detached), so training the draft head never pulls on the policy trunk. +# The default soft cross-entropy gives the forward-KL student gradient +# (``softmax(student) - softmax(teacher)``); the teacher is the policy's own detached next-token +# distribution, so training the draft head never pulls on the policy trunk. The draft-loss design +# is inspired by NeMo-RL (https://github.com/NVIDIA-NeMo/RL), heavily adapted for large-vocab memory. from typing import Optional @@ -24,54 +19,25 @@ def build_teacher_logits( mtp_layer_number: int = 0, detach: bool = True, ) -> torch.Tensor: - """Build the soft-distillation teacher for a given MTP depth. - - The policy's main logits at position ``t`` are the model's distribution over - ``seq[t + 1]``. MTP layer ``k`` (0-indexed) predicts ``seq[t + k + 2]`` at - position ``t``, whose teacher distribution (conditioned on the verified - prefix) is the main model's distribution at position ``t + k + 1`` — i.e. - ``main_logits`` rolled left by ``k + 1``. This generalizes NeMo-RL's - single-step roll (its Eagle draft is one layer, rolled by ``-1``). - - Args: - main_logits: ``[batch, seq, vocab]`` policy logits. - mtp_layer_number: 0-indexed MTP depth ``k``. - detach: Detach the teacher so no gradient flows back into the policy - trunk (the decoupling that makes this "draft" training). + """Build the soft-distillation teacher for MTP depth ``k`` (0-indexed). - Returns: - ``[batch, seq, vocab]`` teacher logits, rolled and (optionally) detached. - The wrapped-around tail positions are invalid; the caller's loss mask - must zero them out (see ``shift_mask_for_mtp``). + Depth ``k`` predicts ``seq[t+k+2]`` at position ``t``, whose teacher distribution is the policy's + own distribution at ``t+k+1`` — i.e. ``main_logits`` rolled left by ``k+1`` (and detached, so no + gradient reaches the trunk). The wrapped tail positions are invalid; the caller's loss mask zeros + them (see ``shift_mask_for_mtp``). """ teacher = main_logits.detach() if detach else main_logits return torch.roll(teacher, shifts=-(mtp_layer_number + 1), dims=1) def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.Tensor: - """Roll a ``[batch, seq]`` loss mask to align with an MTP teacher/label. - - MTP depth ``k`` supervises position ``t`` against a target read at ``t + k + 1`` - (the policy distribution / label for ``seq[t + k + 2]``). For that supervision to - be valid, BOTH the *source* position ``t`` (where the draft head produced its - logits) and the *target* position ``t + k + 1`` (where the teacher/label is read) - must be real tokens. We therefore require ``mask[t] AND mask[t + shift]``: - - * ``rolled = roll(mask, -shift)`` is the target-side validity ``mask[t + shift]``, - with the wrapped tail zeroed (no valid target past the sequence end). - * multiplying by the original ``mask`` re-applies the source-side validity. - - The source-side ``AND`` is what makes this correct under **left padding** (and any - padding layout): rolling a left-padded mask left turns the last pad slot into a - ``1`` (its target is the first real token), which would otherwise pull a de-padded - *zero-logit* (→ uniform) pad position into the loss and inflate it by up to - ``log(V)`` per leaked position. Re-ANDing the source mask drops those positions, so - the loss no longer depends on the de-pad zero-fill happening to be masked out. - - Note: each de-padded row holds a single sequence, so the one wrapped boundary is - handled by zeroing ``rolled``'s tail. If multiple sequences are ever packed into a - single row, per-sequence boundaries would need ``roll_tensor``-style cu_seqlens - handling here. + """Roll a ``[batch, seq]`` loss mask to align with an MTP teacher/label at depth ``k``. + + Supervision at position ``t`` is valid only if both the source ``t`` and the target ``t+shift`` + are real tokens, so we AND ``mask[t]`` with ``roll(mask, -shift)`` (tail zeroed). Re-ANDing the + source mask is what keeps left padding correct: rolling a left-padded mask left would otherwise + unmask a pad slot whose de-padded zero-logit (uniform) inflates the loss by up to ``log(V)``. + Assumes one sequence per row; packing multiple would need cu_seqlens-aware boundaries. """ shift = mtp_layer_number + 1 rolled = torch.roll(mask, shifts=-shift, dims=1) @@ -86,16 +52,11 @@ def unpadded_vocab_shard_width( ) -> int: """Width of this rank's vocab shard excluding Megatron's padded tail. - When the global vocab is padded to divide across TP (``should_pad_vocab``), the padding occupies - the tail of the last rank(s)' shards (the vocab is partitioned contiguously). Padded weight rows - are never trained (zero-filled by the HF conversion), so their logits must not enter the draft - loss: in the teacher softmax they leak ~uniform probability mass, and in the teacher top-k they - can steal slots from real tokens. Callers slice ``logits[..., :width]`` — a view, so autograd - zero-fills the padded tail's gradient and per-rank widths may differ (the loss collectives only - reduce ``[..., 1]``/``[...]``-shaped tensors, never the vocab dim). - - Returns ``vocab_shard_width`` unchanged when ``true_vocab_size`` is ``None`` (unknown) or the - shard holds no padding. + When the vocab is padded to divide across TP (``should_pad_vocab``), the padding lands in the tail + of the last rank(s)' shards. Those rows are never trained, so their logits must not enter the draft + loss (they leak ~uniform mass into the teacher softmax / top-k). Callers slice ``logits[..., :width]`` + (a view; autograd zero-fills the dropped tail's grad). Returns the full width when ``true_vocab_size`` + is ``None`` or the shard holds no padding. """ if true_vocab_size is None: return vocab_shard_width @@ -104,10 +65,8 @@ def unpadded_vocab_shard_width( def _vocab_parallel_softmax(vocab_parallel_logits, group): - """Global softmax over a TP-sharded vocab dim. Allocates a single full-vocab output tensor (the - ``- logits_max`` subtraction) and does the rest in place on it; the input is NOT mutated (so it is - safe even when the caller passes float32 logits, where ``.float()`` is a no-op). Mirrors NeMo-RL's - ``_compute_distributed_softmax``.""" + """Global softmax over a TP-sharded vocab dim. Allocates one full-vocab output (the ``- logits_max`` + subtraction) and does the rest in place on it; the input is not mutated.""" logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) exp_logits = (vocab_parallel_logits - logits_max).exp_() # new buffer, then in place @@ -118,10 +77,8 @@ def _vocab_parallel_softmax(vocab_parallel_logits, group): def _vocab_parallel_log_softmax(vocab_parallel_logits, group): """Global log-softmax over a TP-sharded vocab dim. Uses ``torch.logsumexp`` (a fused - ``[.,.,V]->[.,.,1]`` reduction) so it never materializes a full-vocab ``exp`` temporary -- the key - memory fix at a 248K vocab, where that temp is multiple GiB (NeMo-RL's - ``_compute_distributed_log_softmax`` allocates it via an explicit ``.exp().sum()``). Allocates one - full-vocab output buffer; the input is NOT mutated.""" + ``[.,.,V]->[.,.,1]`` reduction) so it never materializes a full-vocab ``exp`` temporary -- multiple + GiB at a 248K vocab. Allocates one full-vocab output buffer; the input is not mutated.""" logits_max = torch.amax(vocab_parallel_logits, dim=-1, keepdim=True) torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=group) shifted = vocab_parallel_logits - logits_max # new buffer (input untouched); values now <= 0 @@ -132,16 +89,14 @@ def _vocab_parallel_log_softmax(vocab_parallel_logits, group): class _VocabParallelSoftCrossEntropy(torch.autograd.Function): - """Soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` when - ``student``/``teacher`` logits are sharded across the vocab (TP) dim. - - Memory-lean port of NeMo-RL's ``DistributedCrossEntropy`` (``nemo_rl/distributed/model_utils.py``): - the per-token cross-entropy is a single ``einsum`` (no full-vocab product tensor) and the student - log-prob buffer is reused **in place** for the backward softmax, so only two full-vocab tensors - are kept (vs ~6-8 in the naive form -- the difference is several GiB at a 248K vocab). Forward - all-reduces across the TP group so the softmax normalizers and the teacher/student dot are over - the *global* vocab. Backward returns the classic cross-entropy gradient - ``softmax(student) - softmax(teacher)`` (teacher is a detached target; no gradient flows to it). + """Soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` for vocab-(TP)-sharded + logits. + + Memory-lean: the per-token CE is a single ``einsum`` (no full-vocab product tensor) and the student + log-prob buffer is reused in place as the backward softmax, so only two full-vocab tensors are kept + (vs ~6-8 naively -- several GiB at a 248K vocab). Forward all-reduces across the TP group so the + normalizers and the dot are over the global vocab; backward returns ``softmax(student) - + softmax(teacher)`` (teacher detached, no gradient). """ @staticmethod @@ -255,24 +210,15 @@ def per_token(student, teacher): class _VocabParallelTopkSoftCE(torch.autograd.Function): - """Top-k approximation of the soft cross-entropy ``-sum_v softmax(teacher)_v * log_softmax(student)_v`` - that never materializes a full-vocab softmax. - - Only the teacher's top-k tokens are distilled, with teacher and student renormalized over that set. - Memory is ``O(seq * k)`` instead of ``O(seq * vocab)`` -- it fits at a 248K vocab *and* avoids the - allocator fragmentation that the full-vocab / chunked paths cause. - - Parallelism: the vocab is sharded over the tensor-parallel ``group``. We take **each rank's local - top-k** of its shard and reconcile the softmax normalizers and the per-token cross-entropy across - the group with ``all_reduce`` (MAX for the stable-softmax shift, SUM for the denominators and the - loss). No ``all_gather`` / global-index gather is needed because student and teacher share the same - vocab partition, so a rank's local top-k indices address both. This works for ANY tensor-parallel - size and is correct across nodes (the collectives run over the configured ``group``); ``group=None`` - / world-size 1 is the single-device path. The distilled set is the union of the per-rank top-k - (<= ``k * tp_size`` tokens) -- a benign, richer signal at larger TP, not a correctness issue. - - Backward is manual (scatter the ``softmax(student) - softmax(teacher)`` gradient back to this rank's - own top-k columns), so no gradient flows through the collectives. + """Top-k approximation of the soft cross-entropy that never materializes a full-vocab softmax. + + Only the teacher's top-k tokens are distilled (teacher and student renormalized over that set), so + memory is ``O(seq*k)`` instead of ``O(seq*vocab)`` -- fits at a 248K vocab without fragmentation. + Under TP, each rank takes its shard's local top-k and reconciles the normalizers and per-token CE + across the ``group`` (MAX for the stable-softmax shift, SUM for the rest); no all-gather is needed + since student and teacher share the vocab partition. The distilled set is the union of per-rank + top-k (<= ``k*tp_size``) -- a benign, richer signal. Backward scatters the + ``softmax(student) - softmax(teacher)`` gradient to this rank's own top-k columns. """ @staticmethod diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 3816a2b788..5b0566560b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -312,18 +312,11 @@ def forward_backward_mini_batch( """ forward_backward_func = get_forward_backward_func() - # ------------------------------------------------------------------ - # Multi-Token Prediction (MTP) — decoupled draft-head training. - # ------------------------------------------------------------------ - # If the model was built with native MTP heads (policy.megatron_config.mtp_num_layers), - # train them with an explicit, decoupled loss instead of Megatron's in-forward - # process_mtp_loss / MTPLossAutoScaler path. This is model-agnostic: any model that exposes a - # native MTP block works (GPTModel for DeepSeek/GLM/Qwen3-Next, MambaModel for Qwen3.5/ - # NemotronH). We keep the heads running inside the model's forward (so we don't have to - # reconstruct rotary embeddings) but pass NO labels, so process_mtp_loss short-circuits and no - # MTP gradient is coupled onto the trunk. A forward hook captures the MTP block's hidden states - # (with its trunk input optionally detached) and we project + score them ourselves. Only - # active during training (not forward_only / logprob-only passes). + # Multi-Token Prediction (MTP): if the model was built with native MTP heads, train them with + # an explicit decoupled loss instead of Megatron's in-forward process_mtp_loss path. The heads + # still run inside the forward (so we reuse their rotary embeddings) but with NO labels, so + # process_mtp_loss short-circuits and no MTP gradient couples onto the trunk; a forward hook + # captures their hidden states (trunk input optionally detached) for us to score. Training only. model_config = get_model_config(self.actor_module[0]) mtp_enabled = (not forward_only) and bool(getattr(model_config, "mtp_num_layers", None)) mcfg = self.cfg.policy.megatron_config @@ -443,25 +436,19 @@ def loss_func(logits, data): rollout_logprobs=rollout_action_logprobs, ) - # --- Decoupled MTP / draft loss ------------------------------------------------- - # Score the (detached-input) MTP head against the policy's own next-token - # distribution (soft CE) or the ground-truth future tokens (hard CE), over every real - # token. Returns a local masked-mean scalar that is folded into the loss below with the - # same reduction treatment as the KL/entropy aux terms. + # Decoupled MTP / draft loss: score the detached-input MTP head against the policy's own + # next-token distribution (soft CE) or the ground-truth future tokens (hard CE). The local + # masked-mean scalar is folded into the loss below like the KL/entropy aux terms. draft_loss = None student_logits_list = data.get("mtp_student_logits") if mtp_enabled and student_logits_list: draft_mask = data["attention_mask"].to(logits.dtype) # [batch, seq_len] vocab_size_tp = logits.shape[-1] - # Undo the in-place temperature scaling so the teacher is the true policy - # distribution (student logits are produced unscaled). + # Undo the in-place temperature scaling so the teacher is the true policy distribution. teacher_src = logits if temperature == 1.0 else logits * temperature - # Megatron may pad the global vocab to divide across TP (should_pad_vocab); padded - # weight rows are never trained, so their logits would leak probability mass into - # the teacher softmax and could steal teacher top-k slots. Slice this rank's shard - # to its true width (a view; autograd zero-fills the tail's grad). model_config here - # is the bridge provider, whose vocab_size is the true (unpadded) HF vocab; the - # vocab_start_index passed to hard CE keeps the PADDED stride (global column offset). + # Megatron may pad the vocab to divide across TP; padded rows are never trained, so slice + # this rank's shard to its true width to keep them out of the teacher (a view; autograd + # zero-fills the tail). hard CE's vocab_start_index keeps the padded (global) stride. true_shard_width = unpadded_vocab_shard_width( getattr(model_config, "vocab_size", None), vocab_size_tp, tp_rank ) @@ -473,10 +460,8 @@ def loss_func(logits, data): for layer_idx, student_logits in enumerate(student_logits_list): if true_shard_width != vocab_size_tp: student_logits = student_logits[..., :true_shard_width] - # The hard label for position t at depth k is the *token* seq[t+k+2] — one - # position past the soft-CE teacher (a *distribution* read at t+k+1) — so its - # validity mask needs one extra shift. Without it, the last in-range position of - # every row trains against a zeroed/pad label. + # Hard CE's target (token seq[t+k+2]) is one position past the soft-CE teacher + # (a distribution at t+k+1), so its validity mask needs one extra shift. mask_shift = layer_idx + 1 if mtp_loss_type == "hard_ce" else layer_idx layer_mask = shift_mask_for_mtp(draft_mask, mask_shift) if mtp_loss_type == "hard_ce": @@ -493,11 +478,9 @@ def loss_func(logits, data): ) else: if mtp_loss_topk: - # Top-k draft loss: O(seq*k) memory, no full-vocab softmax (fits + avoids - # fragmentation at large vocab). Reconciled across the TP group, so it - # scales to any tensor-parallel size incl. cross-node. Pass the *un-rolled* - # policy logits + roll_shift so top-k runs on the policy's own logits (no - # ~[S, vocab] rolled-teacher copy); only the small [B, S, k] top-k is rolled. + # Top-k draft loss: O(seq*k) memory, no full-vocab softmax. Pass the un-rolled + # policy logits + roll_shift so top-k runs on them directly and only the small + # [B, S, k] result is rolled (avoids a full rolled-teacher copy). per_layer_losses.append( draft_soft_ce_topk( student_logits, @@ -520,8 +503,8 @@ def loss_func(logits, data): ) ) draft_loss = torch.stack(per_layer_losses).mean() - # Drop the dict's reference, this microbatch's autograd graph still holds the - # tensor for its own backward, after which it is freed instead of lingering. + # Drop the dict's reference so the tensor is freed after this microbatch's backward + # (the autograd graph still holds it until then) instead of lingering. del data["mtp_student_logits"] student_logits_list = None @@ -770,18 +753,15 @@ def depad(tensor): "cannot align against)." ) - # Run the policy forward. When MTP is active, a pre-hook records the native MTP block's - # arguments (we pass NO labels, so the model's process_mtp_loss short-circuits and the - # main logits stay coupled to the trunk; MTPLossAutoScaler never runs). + # Run the policy forward; when MTP is active a pre-hook records the MTP block's arguments. student_hidden = None student_model = None with maybe_capture_mtp_hidden( model, mtp_enabled, detach_trunk=mtp_detach_trunk, - # Full shared-weight isolation also requires detaching the MTP block's re-embedding - # of the rolled input ids (the second gradient path into the tied embedding/lm_head, - # next to the output projection handled in project_mtp_hidden_to_logits). + # Full shared-weight isolation also detaches the MTP block's re-embedding of the rolled + # input ids (the second gradient path into the tied embedding/lm_head). detach_shared_embedding=mtp_detach_shared_output, ) as capture: outputs = model( @@ -799,9 +779,8 @@ def depad(tensor): if not self.remove_microbatch_padding: outputs = depad(outputs) - # Project the decoupled MTP hidden states through the shared output layer and de-pad into - # the same [batch, seq_len, vocab/tp] layout as the main logits, so the draft loss can be - # scored against the policy's own distribution (or hard labels) in loss_func. + # Project the MTP hidden states through the shared output layer and de-pad to the main + # logits' [batch, seq_len, vocab/tp] layout, so loss_func can score the draft loss. if student_hidden is not None: student_logits = project_mtp_hidden_to_logits( student_hidden, student_model, detach_output_weight=mtp_detach_shared_output diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 092426fb53..1c5541b0fa 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -223,57 +223,33 @@ class MegatronConfig(BaseConfig): """If True, freeze MoE router parameters so they are not updated during training. No-op on non-MoE models.""" mtp_num_layers: Optional[int] = None - """Number of Multi-Token Prediction (MTP) layers/heads to build on the model. - - - ``None`` (default): honor the model's own HF config (``num_nextn_predict_layers``). For models - that ship MTP heads (e.g. GLM-4.7-Flash / DeepSeek-V3) this enables and trains them; for models - without MTP heads it is a no-op. - - An ``int``: explicitly override (e.g. ``0`` to force-disable MTP even for MTP-capable models). - - When MTP layers are present, SkyRL adds the Megatron MTP loss to the training step so the heads - track the updated policy, and the heads are synced to vLLM for speculative decoding (set - ``generator.inference_engine.speculative_config`` accordingly).""" + """Number of Multi-Token Prediction (MTP) heads to build. ``None`` honors the model's HF config + (``num_nextn_predict_layers``); an int overrides it (``0`` force-disables MTP). Active heads are + trained with the decoupled draft loss and synced to vLLM for speculative decoding.""" mtp_loss_weight: float = 0.1 - """Weight ``w`` of the decoupled MTP/draft loss in ``policy_loss + w * draft_loss``. - - SkyRL trains the native MTP heads with an *explicit* loss on *detached* trunk hidden states (so - the draft gradient never pulls on the policy backbone) and adds it to the policy loss with this - weight, instead of Megatron's opaque in-forward MTP backward-scale. Only used when MTP layers - are active.""" + """Weight ``w`` of the draft loss in ``policy_loss + w * draft_loss``. The draft loss trains the + MTP heads on *detached* trunk hidden states, so it never pulls on the policy backbone. Only used + when MTP heads are active.""" mtp_loss_type: str = "soft_ce" - """Supervision for the MTP/draft head: - - - ``"soft_ce"`` (default): soft cross-entropy distillation against the policy's own detached, - rolled next-token distribution (matches NeMo-RL's draft training). - - ``"hard_ce"``: standard next-token cross-entropy against the ground-truth future tokens - (classic MTP-pretraining objective).""" + """Draft-head supervision: ``"soft_ce"`` distills against the policy's own detached, rolled + next-token distribution; ``"hard_ce"`` is standard next-token cross-entropy.""" mtp_detach_trunk: bool = True - """If True (default, matches NeMo-RL), detach the trunk hidden states feeding the MTP head so - only the MTP-head parameters receive the draft gradient. If False, let the draft gradient flow - back into the policy backbone (closer to Megatron's native coupled MTP).""" + """Detach the trunk hidden states feeding the MTP head so only the head's parameters get the draft + gradient. If False, the gradient flows back into the policy backbone (Megatron's coupled MTP).""" mtp_detach_shared_output: bool = True - """If True (default), fully isolate the shared embedding/output weights from the draft gradient: - the draft-loss projection uses a detached output weight AND the MTP block's re-embedding of the - rolled input ids is detached, so only the MTP-head parameters (enorm/hnorm/eh_proj/layer/final - norm) train. This matters for tied-embedding models (e.g. Qwen3.5), where a trainable shared - head means the draft loss directly nudges the policy's own logits every step. Matches slime's - MTP-RL training (only ``.mtp.`` params receive gradients). Set False for NeMo-RL's behaviour - (shared embedding/head also trained by the draft loss).""" + """Also isolate the shared embedding/output weights from the draft loss (detach both the output + projection and the MTP block's re-embedding), so only the ``.mtp.`` head params train. Matters for + tied-embedding models (e.g. Qwen3.5), where a trainable shared head lets the draft loss nudge the + policy's own logits. Set False to train the shared embedding/head with the draft loss too.""" mtp_loss_chunk_size: Optional[int] = 1024 - """Sequence-chunk size for the decoupled MTP/draft loss. The draft loss materializes full - ``[batch, seq, vocab]`` softmax tensors; at large vocab (e.g. Qwen3.5's 248K) computing the whole - response at once OOMs. Chunking the loss over the sequence (with gradient checkpointing) bounds - peak memory to one chunk's vocab tensors -- numerically identical, only activation memory differs. - ``None`` disables chunking (whole sequence at once). Lower it if the draft loss still OOMs. - Ignored when ``mtp_loss_topk`` is set (top-k never materializes the full vocab, so no chunking).""" + """Sequence-chunk size for the draft loss, with gradient checkpointing, to bound peak memory at + large vocab (e.g. Qwen3.5's 248K, where the full-sequence softmax OOMs). Numerically identical to + no chunking. ``None`` disables it; ignored when ``mtp_loss_topk`` is set.""" mtp_loss_topk: Optional[int] = None - """If set, use a **top-k** approximation of the soft-CE draft loss: distill only the teacher's top-k - tokens (renormalized over that set) instead of the full vocabulary. Memory is ``O(seq*k)`` instead - of ``O(seq*vocab)``, which both fits and avoids allocator fragmentation at large vocab (e.g. - Qwen3.5's 248K) -- where even the chunked full-vocab loss OOMs/fragments. Scales to any - tensor-parallel size (incl. cross-node): the per-rank top-k is reconciled across the TP group with - all-reduce. Only applies to ``mtp_loss_type="soft_ce"``. ``None`` uses the exact full-vocab loss. - Typical: 64-128 (acceptance is governed by the top tokens, so the dropped tail is benign).""" + """If set, use a top-k approximation of the soft-CE draft loss: distill only the teacher's top-k + tokens (renormalized), ``O(seq*k)`` memory instead of ``O(seq*vocab)`` -- fits at large vocab + without fragmentation. Reconciled across the TP group, so it scales to any parallel size. Soft-CE + only; ``None`` uses the exact full-vocab loss. Typical: 64-128.""" def __post_init__(self): # Backfill defaults for any keys the user didn't override so an override dict @@ -617,12 +593,10 @@ class InferenceEngineConfig(BaseConfig): engine_init_kwargs: Dict[str, Any] = field(default_factory=dict) """Pass-through kwargs for the vLLM engine. Names must match the engine's args.""" speculative_config: Optional[Dict[str, Any]] = None - """Speculative-decoding config passed through to vLLM (``AsyncEngineArgs.speculative_config``). - Use this to enable Multi-Token Prediction (MTP) draft decoding for faster rollout, e.g. - ``{"method": "mtp", "num_speculative_tokens": 1}``. With ``method="mtp"`` vLLM loads the MTP - heads from the *same* policy checkpoint, so SkyRL's weight sync keeps the draft in sync with - the trained policy (requires ``policy.megatron_config.mtp_num_layers`` > 0 on the training side - for the heads to actually be trained). ``None`` disables speculative decoding.""" + """Speculative-decoding config passed through to vLLM (``AsyncEngineArgs.speculative_config``), + e.g. ``{"method": "mtp", "num_speculative_tokens": 1}`` for MTP draft decoding. With + ``method="mtp"`` vLLM loads the MTP heads from the policy checkpoint, kept in sync by weight sync + (needs ``policy.megatron_config.mtp_num_layers`` > 0 to train them). ``None`` disables it.""" override_existing_update_group: str = "auto" """``"auto"``, ``"enable"``, or ``"disable"``.""" external_proxy_url: Optional[str] = None @@ -719,27 +693,19 @@ class EnvironmentConfig(BaseConfig): @dataclass class MTPConfig(BaseConfig): - """High-level Multi-Token Prediction (MTP) control. - - This is the single user-facing knob for MTP. When ``enabled``, ``validate_cfg`` propagates it to - the training side (``policy.megatron_config.mtp_*`` — train the model's native MTP heads with a - decoupled draft loss; the head count comes from the checkpoint's own HF config, overridable via - ``policy.megatron_config.mtp_num_layers``) and the inference side - (``generator.inference_engine.speculative_config`` — vLLM MTP speculative decoding with - ``num_speculative_tokens`` draft tokens per step). The trained heads are kept in sync with the - policy via weight sync, so the draft tracks the updated policy. + """High-level, single user-facing knob for Multi-Token Prediction (MTP). - The trunk hidden states and embeddings feeding the heads are always detached (decoupled) — that - is intrinsic to draft/Eagle training, not a tunable, so it is not exposed here.""" + When ``enabled``, ``validate_cfg`` propagates it to the training side + (``policy.megatron_config.mtp_*``) and the inference side + (``generator.inference_engine.speculative_config``), and weight sync keeps the heads tracking the + policy. The trunk hidden states feeding the heads are always detached, so it is not a tunable.""" enabled: bool = False """Whether to train MTP draft heads and use them for speculative decoding.""" num_speculative_tokens: int = 1 - """Draft depth: how many tokens vLLM speculates per decoding step. Decoupled from the number of - trained MTP heads — with a single-head checkpoint (Qwen3.5/Qwen3-Next/DeepSeek-V3 all ship one), - depths > 1 reuse that head autoregressively at draft time (vLLM never indexes beyond head 0). - Expect per-position acceptance to decay with depth, since the head is trained at depth 1 on true - trunk hidden states but consumes its own outputs at deeper draft positions.""" + """Draft depth vLLM speculates per step, independent of the trained head count. Single-head + checkpoints (Qwen3.5/Qwen3-Next/DeepSeek-V3) reuse the one head autoregressively at depths > 1; + expect acceptance to decay with depth (the head is trained at depth 1).""" loss_type: str = "soft_ce" """``"soft_ce"`` (distill against the policy's own next-token distribution) or ``"hard_ce"``.""" loss_weight: float = 0.1 From 8c9d7eb57da22dfb4a30fb0563bb94ffeb778fda Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Thu, 18 Jun 2026 23:35:48 +0000 Subject: [PATCH 20/28] support sequence packing with mtp --- skyrl/backends/skyrl_train/mtp/soft_ce.py | 37 +++++++- .../megatron/megatron_model_wrapper.py | 86 +++++++++++++++---- .../workers/megatron/megatron_worker.py | 17 ++-- 3 files changed, 114 insertions(+), 26 deletions(-) diff --git a/skyrl/backends/skyrl_train/mtp/soft_ce.py b/skyrl/backends/skyrl_train/mtp/soft_ce.py index bceab99693..725e507b0c 100644 --- a/skyrl/backends/skyrl_train/mtp/soft_ce.py +++ b/skyrl/backends/skyrl_train/mtp/soft_ce.py @@ -30,18 +30,47 @@ def build_teacher_logits( return torch.roll(teacher, shifts=-(mtp_layer_number + 1), dims=1) -def shift_mask_for_mtp(mask: torch.Tensor, mtp_layer_number: int = 0) -> torch.Tensor: +def shift_mask_for_mtp( + mask: torch.Tensor, mtp_layer_number: int = 0, cu_seqlens: Optional[torch.Tensor] = None +) -> torch.Tensor: """Roll a ``[batch, seq]`` loss mask to align with an MTP teacher/label at depth ``k``. Supervision at position ``t`` is valid only if both the source ``t`` and the target ``t+shift`` are real tokens, so we AND ``mask[t]`` with ``roll(mask, -shift)`` (tail zeroed). Re-ANDing the source mask is what keeps left padding correct: rolling a left-padded mask left would otherwise unmask a pad slot whose de-padded zero-logit (uniform) inflates the loss by up to ``log(V)``. - Assumes one sequence per row; packing multiple would need cu_seqlens-aware boundaries. + + With THD sample packing (``trainer.remove_microbatch_padding=true``) the whole micro-batch is one + packed row ``[1, T]`` holding many concatenated sub-sequences, so a single ``torch.roll`` would + leak each sub-sequence's target across its boundary into the next one (slime calls the fix the + "roll tensor update for MTP"). Pass ``cu_seqlens`` (the packed *padded* segment boundaries, + ``packed_seq_params.cu_seqlens_q_padded``) to roll *within* each segment and zero the trailing + ``shift`` positions of every segment instead of only the global tail. This masks exactly the + positions whose teacher/label roll crosses a boundary, so those rolls can stay plain global + ``torch.roll`` (their wrong values land only on now-zeroed positions). Mirrors the CP=1 branch of + megatron-core's ``_roll_tensor_packed_seq``; CP>1 is rejected upstream for MTP. """ shift = mtp_layer_number + 1 - rolled = torch.roll(mask, shifts=-shift, dims=1) - rolled[:, -shift:] = 0 + if cu_seqlens is None: + rolled = torch.roll(mask, shifts=-shift, dims=1) + rolled[:, -shift:] = 0 + return mask * rolled + + # Packed: roll each [start, end) segment independently, zeroing its last `shift` positions. + bounds = cu_seqlens.tolist() if torch.is_tensor(cu_seqlens) else list(cu_seqlens) + rolled = torch.zeros_like(mask) + for i in range(len(bounds) - 1): + start, end = int(bounds[i]), int(bounds[i + 1]) + seg_len = end - start + if seg_len <= 0: + continue + seg_rolled = torch.roll(mask[:, start:end], shifts=-shift, dims=1) + # Zero the wrapped tail; the whole segment if it is no longer than the shift. + if shift < seg_len: + seg_rolled[:, -shift:] = 0 + else: + seg_rolled.zero_() + rolled[:, start:end] = seg_rolled return mask * rolled diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 5b0566560b..a239ee31eb 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -90,6 +90,40 @@ def _build_packed_targets( return targets.unsqueeze(0) +def _build_packed_valid_mask( + attention_mask: torch.Tensor, + packed_seq_params, + sub_seq_lengths: Optional[list[list[int]]] = None, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Build a ``[1, T]`` real-token mask aligned to the packed (THD) logits layout. + + 1.0 for real tokens, 0.0 for the per-segment alignment padding that ``preprocess_packed_seqs`` + inserts between sub-sequences. This is the packed counterpart of the ``[batch, seq]`` + ``attention_mask`` the decoupled MTP draft loss uses to mask invalid positions; mirrors the + index math of :func:`_build_packed_targets` but scatters ones instead of token ids. + """ + cu_padded = packed_seq_params.cu_seqlens_q_padded.to(device=attention_mask.device, dtype=torch.long) + total_padded_tokens = int(cu_padded[-1].item()) + mask = torch.zeros((total_padded_tokens,), dtype=dtype, device=attention_mask.device) + if sub_seq_lengths is not None: + cu_padded_cpu = cu_padded.detach().cpu().tolist() + seg_idx = 0 + for row_lens in sub_seq_lengths: + for seq_len in row_lens: + seq_len = int(seq_len) + packed_start = cu_padded_cpu[seg_idx] + mask[packed_start : packed_start + seq_len] = 1.0 + seg_idx += 1 + return mask.unsqueeze(0) + + attn = attention_mask.to(device=cu_padded.device, dtype=torch.bool) + token_offsets = attn.to(torch.long).cumsum(dim=1) - 1 + packed_indices = cu_padded[:-1].unsqueeze(1) + token_offsets + mask[packed_indices[attn]] = 1.0 + return mask.unsqueeze(0) + + class MegatronModelWrapper: def __init__( self, @@ -442,7 +476,13 @@ def loss_func(logits, data): draft_loss = None student_logits_list = data.get("mtp_student_logits") if mtp_enabled and student_logits_list: - draft_mask = data["attention_mask"].to(logits.dtype) # [batch, seq_len] + # Under THD sample packing the teacher (logits) and student logits are packed + # ([1, T, vocab]); use the matching packed real-token mask + segment boundaries so the + # MTP roll/mask respect sub-sequence boundaries (see shift_mask_for_mtp). Otherwise the + # [batch, seq] attention mask aligns with the de-padded logits. + packed = packed_seq_params is not None + draft_mask = (data["mtp_packed_mask"] if packed else data["attention_mask"]).to(logits.dtype) + mtp_cu_seqlens = packed_seq_params.cu_seqlens_q_padded if packed else None vocab_size_tp = logits.shape[-1] # Undo the in-place temperature scaling so the teacher is the true policy distribution. teacher_src = logits if temperature == 1.0 else logits * temperature @@ -454,7 +494,10 @@ def loss_func(logits, data): ) if true_shard_width != vocab_size_tp: teacher_src = teacher_src[..., :true_shard_width] - hard_labels = build_mtp_next_token_labels(sequences) if mtp_loss_type == "hard_ce" else None + # Packed: build labels from packed_targets ([1, T]) so they align with the packed + # student/teacher; the global roll's cross-segment entries are zeroed by layer_mask. + hard_label_src = data["packed_targets"] if packed else sequences + hard_labels = build_mtp_next_token_labels(hard_label_src) if mtp_loss_type == "hard_ce" else None per_layer_losses = [] for layer_idx, student_logits in enumerate(student_logits_list): @@ -463,7 +506,7 @@ def loss_func(logits, data): # Hard CE's target (token seq[t+k+2]) is one position past the soft-CE teacher # (a distribution at t+k+1), so its validity mask needs one extra shift. mask_shift = layer_idx + 1 if mtp_loss_type == "hard_ce" else layer_idx - layer_mask = shift_mask_for_mtp(draft_mask, mask_shift) + layer_mask = shift_mask_for_mtp(draft_mask, mask_shift, cu_seqlens=mtp_cu_seqlens) if mtp_loss_type == "hard_ce": layer_labels = torch.roll(hard_labels, shifts=-(layer_idx + 1), dims=1) per_layer_losses.append( @@ -732,9 +775,10 @@ def forward_step(batch_iter, model): # Recover [batch, seq_len, ...] from Megatron's internal (left-removed) layout. Only used # on the non-packed path: with sample packing (remove_microbatch_padding) the logits stay - # packed and loss_func consumes packed_targets instead. MTP draft training requires the - # non-packed path (the teacher main-logits and the de-padded student logits must share the - # [batch, seq, vocab] layout), enforced by the assertion below. + # packed ([1, T, vocab]) and loss_func consumes packed_targets instead. MTP draft training + # keeps the student logits in whichever layout the teacher (main logits) uses — de-padded + # [batch, seq, vocab] without packing, packed [1, T, vocab] with it — so the two always + # align (see the packed-aware mask in loss_func / mtp/soft_ce.py). def depad(tensor): return recover_left_padding( tensor, @@ -744,14 +788,14 @@ def depad(tensor): post_process=is_last_stage, ) - # MTP draft training de-pads the student logits to [batch, seq, vocab] to align with the - # teacher (main logits). Sample packing keeps the main logits packed (loss_func consumes - # packed_targets), so the two layouts would mismatch — disallow the combination loudly. - assert not (mtp_enabled and self.remove_microbatch_padding), ( - "MTP/draft training requires trainer.remove_microbatch_padding=False " - "(sample packing keeps the policy logits packed, which the decoupled draft loss " - "cannot align against)." - ) + # The decoupled MTP teacher/label rolls are context-parallel-unaware (a plain global + # torch.roll, both here and in the packed path), so MTP draft training requires CP=1. + # CP>1 would need the cross-rank boundary exchange that megatron's roll_tensor implements. + if mtp_enabled: + assert mpu.get_context_parallel_world_size() == 1, ( + "MTP/draft training does not support context parallelism " + "(context_parallel_size > 1): the teacher/label roll is CP-unaware." + ) # Run the policy forward; when MTP is active a pre-hook records the MTP block's arguments. student_hidden = None @@ -779,13 +823,21 @@ def depad(tensor): if not self.remove_microbatch_padding: outputs = depad(outputs) - # Project the MTP hidden states through the shared output layer and de-pad to the main - # logits' [batch, seq_len, vocab/tp] layout, so loss_func can score the draft loss. + # Project the MTP hidden states through the shared output layer so loss_func can score the + # draft loss. Match the teacher's layout: keep them packed ([1, T, vocab/tp]) under sample + # packing, or de-pad to [batch, seq_len, vocab/tp] otherwise. When packed, also hand + # loss_func the packed real-token mask so it can mask cross-segment MTP targets. if student_hidden is not None: student_logits = project_mtp_hidden_to_logits( student_hidden, student_model, detach_output_weight=mtp_detach_shared_output ) - batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] + if self.remove_microbatch_padding: + batch["mtp_student_logits"] = student_logits + batch["mtp_packed_mask"] = _build_packed_valid_mask( + attention_mask, packed_seq_params, sub_seq_lengths=sub_seq_lengths + ) + else: + batch["mtp_student_logits"] = [depad(sl) for sl in student_logits] if rollout_expert_indices is not None: setup_per_microbatch_replay_backward() diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 20a6a458ca..496681a0c2 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -379,11 +379,18 @@ def init_configs( provider = bridge.to_megatron_provider() - # Disable MTP for training: its aux loss is unused, and under full - # recompute its checkpointed forward passes packed_seq_params positionally - # into tensor_parallel.checkpoint (tensors only), breaking packed-sequence - # backward. Mirrors the MTP-disable in model_bridges.py. - if getattr(provider, "mtp_num_layers", None): + # Disable MTP for ordinary training: Megatron's in-forward MTP aux loss is unused, and under + # full recompute its checkpointed forward passes packed_seq_params positionally into + # tensor_parallel.checkpoint (tensors only), breaking packed-sequence backward. + # EXCEPTION: decoupled MTP/draft training (trainer.mtp.enabled on the policy worker) MUST keep + # the heads built — it runs them in-forward in eval mode (which skips the recompute path that + # breaks, see mtp/hidden_capture.py), scores a separate draft loss via a forward hook, and + # weight-syncs the heads to vLLM. It is compatible with sample packing (the draft loss masks + # cross-segment positions, see mtp/soft_ce.py::shift_mask_for_mtp), so we do NOT force packing + # off here. The MTP-config block below resolves the final head count. + _mtp_cfg = getattr(self.cfg, "mtp", None) + _mtp_training = enable_mtp and _mtp_cfg is not None and getattr(_mtp_cfg, "enabled", False) + if not _mtp_training and getattr(provider, "mtp_num_layers", None): logger.info(f"Disabling MTP for training (mtp_num_layers={provider.mtp_num_layers} -> None)") provider.mtp_num_layers = None From 5c477db856393780b579b206b767534a4c1f2100 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 24 Jun 2026 01:37:17 +0000 Subject: [PATCH 21/28] [mtp] Add full separate MTP optimizer + isolated DDP grad buffer (C-full) Isolate the MTP/draft head into its own Megatron DDP grad buffer + DistributedOptimizer so the policy's distributed grad reduction is byte-identical to a no-MTP model, while the head co-trains at full strength: - mtp/separate_optim.py: pre-wrap freeze, SeparateMTPOptimizer (own DDP buffer + optimizer + scheduler), policy-finalize exclusion, hidden() to exclude the head from the policy grad-norm/clip. - megatron_worker.py: wire C-full (freeze pre-wrap -> separate optimizer -> per-iter buffer zero -> separate grad-norm/clip + step), plus checkpoint sidecar save/load for the head optimizer state. - config: mtp_separate_optimizer toggle (MegatronConfig). - trainer.py: guard the rollout/train logprob-diff metric against an empty loss_mask batch. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../skyrl_train/mtp/separate_optim.py | 239 ++++++++++++++++++ .../workers/megatron/megatron_worker.py | 104 +++++++- skyrl/train/config/config.py | 9 + skyrl/train/trainer.py | 28 +- 4 files changed, 367 insertions(+), 13 deletions(-) create mode 100644 skyrl/backends/skyrl_train/mtp/separate_optim.py diff --git a/skyrl/backends/skyrl_train/mtp/separate_optim.py b/skyrl/backends/skyrl_train/mtp/separate_optim.py new file mode 100644 index 0000000000..c1c45a3ef7 --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/separate_optim.py @@ -0,0 +1,239 @@ +"""C-full: isolate the MTP / draft head into its OWN grad buffer + ``DistributedOptimizer``. + +Why this exists +--------------- +The decoupled draft loss is autograd-clean (it only reaches ``.mtp.*`` params; the trunk hidden and +the shared output/embedding weights are detached). But when the MTP head shares the policy's single +Megatron DDP grad buffer, the head's mere *presence* changes the layout of that buffer and therefore +the floating-point result of the POLICY gradient's distributed reduction (the DP reduce-scatter and the +TP all-reduce of sequence-parallel / layernorm grads, both of which coalesce every ``requires_grad`` +param into one flat tensor). That perturbation is deterministic and systematic, and the RL feedback +loop amplifies it into the policy-entropy collapse observed on MiMo-7B with MTP enabled. + +Neither separating the grad *clip* (the head's grad does not dilute the policy via the clip — verified) +nor pinning the NCCL algorithm fixes it, because the channel is the shared *buffer/reduction* itself. +The robust fix is to make the policy's grad buffer + reduction byte-identical to a model built with NO +MTP head, while still co-training the head at full strength. + +Mechanism +--------- +Megatron exposes no "exclude these params from the buffer" API; the only lever is ``requires_grad`` at +DDP-construction time (DDP buckets only ``requires_grad`` params, and changing it *after* the wrap +deadlocks distributed-optimizer init). So: + +1. :func:`freeze_mtp_params_pre_wrap` (a provider pre-wrap hook) sets the ``.mtp`` params + ``requires_grad=False`` *before* the policy is DDP-wrapped, so they are excluded from the policy + grad buffer -> the policy buffer is byte-identical to the no-MTP build. +2. After the policy optimizer is built, :class:`SeparateMTPOptimizer` flips the head back to + ``requires_grad=True`` and wraps the SAME ``host.mtp`` submodule (left in place inside ``GPTModel``, + so capture / replay / weight-export are unchanged) in its own Megatron DDP + ``DistributedOptimizer``. +3. The policy's own ``finalize_model_grads`` would still coalesce the head's layernorm grads into the + policy TP all-reduce (it iterates every ``requires_grad`` param of the GPTModel). :func:`make_policy_finalize_excluding_mtp` + wraps it to transiently hide the head during the policy finalize (backward is already complete, so + this is only a read-time filter), keeping the policy reduction byte-identical to the no-MTP build. +4. The combined ``policy + weight * draft`` backward auto-routes each param's grad to its own buffer + (the draft loss is autograd-decoupled). We finalize + step the head's optimizer alongside the policy's. + +Acceptance test: with C-full on, the fixed-batch no-clip policy-param trajectory must match the no-MTP +run within the run-to-run nondeterminism floor (see +``tests/.../megatron/test_mtp_grad_coupling.py::test_mtp_fixed_batch_param_drift``). +""" + +from __future__ import annotations + +import contextlib +from typing import List, Optional, Union + +import torch +from megatron.core import parallel_state as mpu +from megatron.core.distributed import DistributedDataParallel +from megatron.core.distributed import ( + finalize_model_grads as _mcore_finalize_model_grads, +) +from megatron.core.distributed.distributed_data_parallel_config import ( + DistributedDataParallelConfig, +) +from megatron.core.utils import get_model_config + +from skyrl.train.config.config import get_config_as_dict + + +def is_mtp_param_name(name: str) -> bool: + """True for parameter names that belong to the MTP / draft head.""" + return ".mtp." in name or name.startswith("mtp.") + + +def _resolve_mtp_module(policy_module): + """Return the ``host.mtp`` submodule (or None) for a (possibly DDP/Float16-wrapped) policy chunk.""" + from skyrl.backends.skyrl_train.mtp.hidden_capture import ( + _resolve_mtp_host, + _unwrap_model, + ) + + host = _resolve_mtp_host(_unwrap_model(policy_module)) + return getattr(host, "mtp", None) + + +def freeze_mtp_params_pre_wrap(model_or_models: Union[torch.nn.Module, List[torch.nn.Module]]): + """Provider pre-wrap hook: set the MTP/draft-head params ``requires_grad=False`` so Megatron's DDP + excludes them from the POLICY grad buffer (the only Megatron lever for buffer exclusion is + ``requires_grad`` at construction time). They are re-enabled and given their own buffer + optimizer + by :class:`SeparateMTPOptimizer` after the policy is wrapped. Result: the policy grad buffer and its + distributed reduction are byte-identical to a model built with no MTP head. + + Modifies in place AND returns the model unchanged. (The bridge's ``pre_wrap_hook`` property + composes hooks via ``model = hook(model)`` with NO None-guard, so a hook that returns None would + break the chain for any subsequent hook — we must return the model.) + """ + models = model_or_models if isinstance(model_or_models, list) else [model_or_models] + for model in models: + for name, param in model.named_parameters(): + if is_mtp_param_name(name): + param.requires_grad = False + return model_or_models + + +def make_policy_finalize_excluding_mtp(mtp_params: List[torch.nn.Parameter]): + """Wrap Megatron's ``finalize_model_grads`` so the MTP head is hidden during the POLICY finalize. + + The policy finalize iterates every ``requires_grad`` param of the GPTModel (the head lives inside it) + and coalesces their sequence-parallel/layernorm grads into ONE tensor-parallel all-reduce. If the + head's grads were included, that coalesced reduction would differ from the no-MTP build -> the exact + perturbation we are eliminating. Backward is already complete when finalize runs, so transiently + clearing the head's ``requires_grad`` is a pure read-time filter (no effect on the head's grads, + which already live in the separate MTP buffer and are reduced by :meth:`SeparateMTPOptimizer.step`). + """ + + def finalize(model, *args, **kwargs): + saved = [(p, p.requires_grad) for p in mtp_params] + for p in mtp_params: + p.requires_grad = False + try: + return _mcore_finalize_model_grads(model, *args, **kwargs) + finally: + for p, rg in saved: + p.requires_grad = rg + + return finalize + + +def _build_mtp_ddp_config(ddp_config) -> DistributedDataParallelConfig: + """Build the head's ``DistributedDataParallelConfig`` from the policy's, but with overlap disabled. + + ``overlap_grad_reduce=False`` so the head's grads simply accumulate into its buffer across + micro-batches and are reduced once in :meth:`SeparateMTPOptimizer.finalize_grads` — there is no + per-microbatch sync to coordinate with the policy's pipeline schedule (the head is not in the chunk + list passed to ``forward_backward_func``). ``overlap_param_gather=False`` for the same reason. + """ + cfg = DistributedDataParallelConfig() + cfg.use_distributed_optimizer = True + if ddp_config is not None: + for k, v in get_config_as_dict(ddp_config).items(): + setattr(cfg, k, v) + cfg.overlap_grad_reduce = False + cfg.overlap_param_gather = False + return cfg + + +class SeparateMTPOptimizer: + """Owns the MTP/draft head's isolated grad buffer + ``DistributedOptimizer`` (C-full). + + The head stays inside the policy ``GPTModel`` (so capture / replay / weight-export are unchanged) but + its params are excluded from the policy DDP buffer (via :func:`freeze_mtp_params_pre_wrap`) and wrapped + here in a SEPARATE Megatron DDP, so the policy's distributed grad reduction is byte-identical to a + no-MTP model. The head is co-trained at full strength under this optimizer. + """ + + def __init__(self, policy_module, ddp_config, optim_config, scheduler_config, num_training_steps: int): + from skyrl.backends.skyrl_train.distributed.megatron.optimizer import ( + get_megatron_optimizer, + get_megatron_optimizer_param_scheduler, + ) + + self.mtp_module = _resolve_mtp_module(policy_module) + if self.mtp_module is None: + raise RuntimeError( + "SeparateMTPOptimizer: no `.mtp` head found on the policy model. Enable MTP " + "(trainer.mtp.enabled / mtp_num_layers) or disable mtp_separate_optimizer." + ) + + # Re-enable grads on the head (frozen during the policy wrap by the pre-wrap hook) BEFORE + # building its DDP + optimizer. Setting requires_grad here — before this fresh wrap and its + # optimizer init — avoids the post-wrap requires_grad change that deadlocks dist-optimizer init. + for p in self.mtp_module.parameters(): + p.requires_grad = True + self.mtp_params: List[torch.nn.Parameter] = list(self.mtp_module.parameters()) + + self._model_config = get_model_config(policy_module) + self.mtp_ddp = DistributedDataParallel( + config=self._model_config, + ddp_config=_build_mtp_ddp_config(ddp_config), + module=self.mtp_module, + ) + self.optimizer = get_megatron_optimizer([self.mtp_ddp], optim_config) + self.scheduler = get_megatron_optimizer_param_scheduler( + optimizer=self.optimizer, + config=scheduler_config, + num_training_steps=num_training_steps, + ) + + def zero_grad_buffer(self) -> None: + """Zero the head's grad buffer at the start of a forward-backward (mirrors the policy chunks).""" + self.mtp_ddp.zero_grad_buffer() + + @contextlib.contextmanager + def hidden(self): + """Temporarily set the head's params ``requires_grad=False`` so the POLICY optimizer's + grad-norm + clip EXCLUDE the head. Megatron computes the policy grad-norm by iterating the + GPTModel's ``requires_grad`` params and reading ``main_grad`` — and the head is *structurally* + still inside the GPTModel (so capture/replay/export work), so without this it would pick up the + head's grad (from its separate buffer) and dilute the policy clip exactly like the un-isolated + run. This is a pure read-time filter while the policy step runs; the head's grads stay in its + own buffer and are clipped/stepped separately by :meth:`step`.""" + saved = [(p, p.requires_grad) for p in self.mtp_params] + for p in self.mtp_params: + p.requires_grad = False + try: + yield + finally: + for p, rg in saved: + p.requires_grad = rg + + def finalize_grads(self) -> None: + """Reduce the head's accumulated grads: DP reduce-scatter (own buffer) + TP all-reduce of the + head's sequence-parallel / layernorm grads. We call the targeted primitives rather than the full + ``finalize_model_grads`` so we never touch the policy buffer and avoid the embedding/PP-cross-stage + logic that assumes a complete GPTModel chunk (it would break for tied-embedding models).""" + from megatron.core.distributed.finalize_model_grads import ( + _allreduce_non_tensor_model_parallel_grads, + ) + + self.mtp_ddp.finish_grad_sync() + _allreduce_non_tensor_model_parallel_grads( + [self.mtp_ddp], self._model_config, mpu.get_tensor_model_parallel_group() + ) + + def step(self) -> Optional[float]: + """Finalize the head's grads, step its optimizer + scheduler, and zero its grads. + + Returns the head's grad norm (after the optimizer's own clip), or None if unavailable. + """ + self.finalize_grads() + _, grad_norm, _ = self.optimizer.step() + self.scheduler.step(1) + self.optimizer.zero_grad() + if grad_norm is not None and hasattr(grad_norm, "item"): + grad_norm = grad_norm.item() + return grad_norm + + # -- checkpointing ------------------------------------------------------------------------------- + def state_dict(self): + return {"optimizer": self.optimizer.state_dict(), "scheduler": self.scheduler.state_dict()} + + def load_state_dict(self, state_dict) -> None: + if not state_dict: + return + if state_dict.get("optimizer") is not None: + self.optimizer.load_state_dict(state_dict["optimizer"]) + if state_dict.get("scheduler") is not None: + self.scheduler.load_state_dict(state_dict["scheduler"]) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 496681a0c2..c011e37186 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -848,6 +848,25 @@ def init_model(self, model_path, num_training_steps: int = 1e9): logger.info("freeze_moe_router=True: freezing MoE router params") self.provider.register_pre_wrap_hook(freeze_moe_router) + # MTP C-full: isolate the draft head into its own grad buffer + optimizer so the policy's + # distributed grad reduction is byte-identical to a no-MTP model (the head's presence in the + # shared grad buffer otherwise perturbs the policy reduction and collapses entropy — see + # mtp/separate_optim.py). Step 1, here: freeze the head BEFORE the policy DDP wrap so Megatron's + # DDP (which only buckets requires_grad params) excludes it from the policy grad buffer. The head + # is re-enabled and given its own buffer + optimizer after the policy optimizer is built. + self._mtp_separate = None + self._mtp_cfull_enabled = bool( + self.cfg.policy.megatron_config.mtp_separate_optimizer and getattr(self.provider, "mtp_num_layers", None) + ) + if self._mtp_cfull_enabled: + from skyrl.backends.skyrl_train.mtp.separate_optim import ( + freeze_mtp_params_pre_wrap, + ) + + self.provider.register_pre_wrap_hook(freeze_mtp_params_pre_wrap) + if self._rank == 0: + logger.info("MTP C-full: registered pre-wrap hook to isolate the draft head's grad buffer") + # wrap with DDP for training self.actor_module = self.make_megatron_module( wrap_with_ddp=True, @@ -889,6 +908,30 @@ def init_model(self, model_path, num_training_steps: int = 1e9): num_training_steps=num_training_steps, ) + # MTP C-full step 2: re-enable the draft head and give it its OWN grad buffer + + # DistributedOptimizer (the policy optimizer above already excludes it — it was frozen + # before the policy wrap). The head co-trains at full strength while the policy reduction + # stays byte-identical to a no-MTP model. See mtp/separate_optim.py. + if self._mtp_cfull_enabled: + from skyrl.backends.skyrl_train.mtp.separate_optim import ( + SeparateMTPOptimizer, + ) + + # Fresh optim config for the head (don't share the policy's object). + mtp_optim_config = init_megatron_optim_config( + self.cfg.policy.optimizer_config, self.cfg.policy.megatron_config.optimizer_config_kwargs + ) + self._mtp_separate = SeparateMTPOptimizer( + policy_module=self.actor_module[0], + ddp_config=self.cfg.policy.megatron_config.ddp_config, + optim_config=mtp_optim_config, + scheduler_config=self.cfg.policy.optimizer_config, + num_training_steps=num_training_steps, + ) + if self._rank == 0: + n = sum(p.numel() for p in self._mtp_separate.mtp_params) + logger.info(f"MTP C-full: isolated draft head ({n:,} params) into its own grad buffer + optimizer") + # create worker model self.model = MegatronModelWrapper( config=self.cfg, @@ -897,6 +940,20 @@ def init_model(self, model_path, num_training_steps: int = 1e9): policy_loss_fn=self.policy_loss_fn, ) + # MTP C-full step 3: hide the draft head from the POLICY finalize (which would otherwise coalesce + # the head's layernorm grads into the policy TP all-reduce, perturbing it). Must run AFTER the + # wrapper, which sets finalize_model_grads_func on this same model_config object. + if self._mtp_separate is not None: + from megatron.core.utils import get_model_config + + from skyrl.backends.skyrl_train.mtp.separate_optim import ( + make_policy_finalize_excluding_mtp, + ) + + get_model_config(self.actor_module[0]).finalize_model_grads_func = make_policy_finalize_excluding_mtp( + self._mtp_separate.mtp_params + ) + self.empty_cuda_cache = self.cfg.policy.megatron_config.empty_cuda_cache # Enable expandable_segments after init so model weights stay in IPC-compatible @@ -1037,6 +1094,9 @@ def forward_backward( for chunk in self.actor_module: # if use distributed optimizer, zero grad buffer will be handled by optimizer chunk.zero_grad_buffer() + # MTP C-full: zero the isolated draft-head grad buffer too (it lives outside self.actor_module). + if self._mtp_separate is not None: + self._mtp_separate.zero_grad_buffer() all_metrics = defaultdict(list) @@ -1213,7 +1273,17 @@ def optim_step(self) -> Optional[float]: """ if self.optimizer is None: raise RuntimeError("optim_step called but policy.inference_only_init=True (no optimizer constructed)") - grad_norm = self.strategy.optimizer_step(self.optimizer, self.model, self.scheduler, name="actor") + + # MTP C-full: hide the draft head during the POLICY step so its grad-norm + clip EXCLUDE the + # head (Megatron's grad-norm iterates the GPTModel's requires_grad params and reads main_grad; + # the head is structurally still in the GPTModel, so without this it dilutes the policy clip). + # The policy step then matches a no-MTP run; the head is finalized + stepped separately after. + if self._mtp_separate is not None: + with self._mtp_separate.hidden(): + grad_norm = self.strategy.optimizer_step(self.optimizer, self.model, self.scheduler, name="actor") + self._mtp_grad_norm = self._mtp_separate.step() + else: + grad_norm = self.strategy.optimizer_step(self.optimizer, self.model, self.scheduler, name="actor") # Reset counter for next accumulation cycle self._micro_batches_accumulated = 0 @@ -1222,6 +1292,38 @@ def optim_step(self) -> Optional[float]: grad_norm = grad_norm.detach().cpu().item() if hasattr(grad_norm, "item") else grad_norm return grad_norm + def save_checkpoint(self, ckpt_dir: str, tokenizer=None): + """Save the policy checkpoint, plus the C-full MTP head's separate optimizer state. + + The isolated draft-head optimizer (``self._mtp_separate``) lives OUTSIDE ``self.optimizer``, so + the strategy's checkpoint does not cover it. Persist its (DP-sharded) state per global rank + alongside the policy checkpoint so resume restores the head's optimizer momentum / scheduler. + """ + super().save_checkpoint(ckpt_dir, tokenizer=tokenizer) + if self._mtp_separate is not None: + os.makedirs(ckpt_dir, exist_ok=True) + torch.save(self._mtp_separate.state_dict(), os.path.join(ckpt_dir, f"mtp_optim_rank{self._rank}.pt")) + + def load_checkpoint(self, ckpt_dir: str, load_optimizer_states: bool = True, load_lr_scheduler_states: bool = True): + states = super().load_checkpoint( + ckpt_dir, + load_optimizer_states=load_optimizer_states, + load_lr_scheduler_states=load_lr_scheduler_states, + ) + # Restore the C-full MTP head's separate optimizer state. Guarded on existence so a checkpoint + # written without C-full (or a topology change) degrades to "fresh head optimizer" instead of + # crashing the resume. + if self._mtp_separate is not None and load_optimizer_states: + mtp_path = os.path.join(ckpt_dir, f"mtp_optim_rank{self._rank}.pt") + if os.path.exists(mtp_path): + self._mtp_separate.load_state_dict(torch.load(mtp_path, map_location="cpu")) + elif self._rank == 0: + logger.warning( + f"MTP C-full: no separate-optimizer checkpoint at {mtp_path}; " + "the draft head's optimizer state was not restored." + ) + return states + def get_lr(self) -> Optional[float]: """ Get current learning rate from optimizer. diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 1c5541b0fa..3984fa9efd 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -241,6 +241,15 @@ class MegatronConfig(BaseConfig): projection and the MTP block's re-embedding), so only the ``.mtp.`` head params train. Matters for tied-embedding models (e.g. Qwen3.5), where a trainable shared head lets the draft loss nudge the policy's own logits. Set False to train the shared embedding/head with the draft loss too.""" + mtp_separate_optimizer: bool = True + """C-full: give the MTP/draft head its OWN grad buffer + DistributedOptimizer, fully isolated from + the policy's. The decoupled draft loss is autograd-clean, but when the head shares the policy's + Megatron DDP grad buffer its mere presence changes the layout — and thus the floating-point result — + of the policy gradient's distributed reduction (DP reduce-scatter + TP all-reduce of layernorm + grads). That perturbation is deterministic and the RL feedback loop amplifies it into policy-entropy + collapse. With this on, the policy's grad buffer and reduction are byte-identical to a model built + with no MTP head, while the head still co-trains at full strength. Only used when MTP heads are + active. See skyrl/backends/skyrl_train/mtp/separate_optim.py.""" mtp_loss_chunk_size: Optional[int] = 1024 """Sequence-chunk size for the draft loss, with gradient checkpointing, to bound peak memory at large vocab (e.g. Qwen3.5's 248K, where the full-sequence softmax OOMs). Numerically identical to diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 44080907c2..12e1979f05 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -1266,18 +1266,22 @@ def fwd_logprobs_values_reward( - action_log_probs[training_input["loss_mask"] > 0] ).abs() - logprobs_diff_max = logprobs_diff.max().item() - logprobs_diff_min = logprobs_diff.min().item() - logprobs_diff_mean = logprobs_diff.mean().item() - logprobs_diff_std = logprobs_diff.std().item() - self.all_metrics.update( - { - "policy/rollout_train_logprobs_abs_diff_max": logprobs_diff_max, - "policy/rollout_train_logprobs_abs_diff_min": logprobs_diff_min, - "policy/rollout_train_logprobs_abs_diff_mean": logprobs_diff_mean, - "policy/rollout_train_logprobs_abs_diff_std": logprobs_diff_std, - } - ) + # Guard: a batch with no trainable response tokens (loss_mask all zero, e.g. every + # response dropped by overlong filtering) leaves logprobs_diff empty, and .max()/.min() + # on a 0-element tensor raises. Skip the diagnostic metrics in that case. + if logprobs_diff.numel() > 0: + logprobs_diff_max = logprobs_diff.max().item() + logprobs_diff_min = logprobs_diff.min().item() + logprobs_diff_mean = logprobs_diff.mean().item() + logprobs_diff_std = logprobs_diff.std().item() + self.all_metrics.update( + { + "policy/rollout_train_logprobs_abs_diff_max": logprobs_diff_max, + "policy/rollout_train_logprobs_abs_diff_min": logprobs_diff_min, + "policy/rollout_train_logprobs_abs_diff_mean": logprobs_diff_mean, + "policy/rollout_train_logprobs_abs_diff_std": logprobs_diff_std, + } + ) return training_input def apply_reward_kl_penalty( From 57f14a043060010f97fa5e28f276e03b63cb260c Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 24 Jun 2026 03:51:02 +0000 Subject: [PATCH 22/28] Move train/eval rollout timing into VLLMMetricsScraper window API Replace the brittle mark_pre_eval()/sample_split() approach (which relied on the caller's generate/eval ordering and threaded generation-time values back in) with an explicit start()/pause()/resume()/stop() window that owns its own timing. The scraper accumulates only un-paused active time, so the trainer and evaluate() just bracket their generation calls. - start(label)/stop() return metrics nested under {label}/ (vllm/train/*, vllm/eval/*); pause/resume gate the throughput denominator. - trainer: wrap train generation (handles dynamic sampling's repeated generates) and eval generation; evaluate()/evaluate_step_wise() resume/pause around each eval generation. - fully-async sample() path unchanged. - Rewrite split tests for the window API. --- skyrl/train/evaluate.py | 21 ++- skyrl/train/trainer.py | 48 ++++-- skyrl/train/utils/vllm_metrics_scraper.py | 111 +++++++++----- tests/train/test_vllm_metrics_scraper.py | 170 ++++++++++++++-------- 4 files changed, 243 insertions(+), 107 deletions(-) diff --git a/skyrl/train/evaluate.py b/skyrl/train/evaluate.py index e0dfc5ad94..17bf6dad05 100644 --- a/skyrl/train/evaluate.py +++ b/skyrl/train/evaluate.py @@ -1,7 +1,7 @@ import time from collections import defaultdict from pathlib import Path -from typing import Any, Dict, List +from typing import TYPE_CHECKING, Any, Dict, List, Optional import torch from loguru import logger @@ -30,6 +30,9 @@ validate_generator_output, ) +if TYPE_CHECKING: + from skyrl.train.utils.vllm_metrics_scraper import VLLMMetricsScraper + @torch.no_grad() async def evaluate( @@ -38,6 +41,7 @@ async def evaluate( cfg: SkyRLTrainConfig, global_step: int | None, tokenizer: AutoTokenizer, + vllm_metrics_scraper: Optional["VLLMMetricsScraper"] = None, ) -> Dict[str, float]: """Runs generation and evaluation of trajectories. @@ -48,6 +52,9 @@ async def evaluate( global_step (int | None): current global step, or `None` to indicate a non-training context (e.g., eval-only) tokenizer (AutoTokenizer): tokenizer to use + vllm_metrics_scraper: when set, the open ``vllm/eval`` window is resumed + around each generation and paused after, so only generation time + counts toward eval throughput. Returns: Dict[str, float]: evaluation metrics @@ -72,7 +79,11 @@ async def evaluate( global_step, ) gen_start = time.monotonic() + if vllm_metrics_scraper is not None: + vllm_metrics_scraper.resume() generator_output: GeneratorOutput = await generator.generate(generator_input) + if vllm_metrics_scraper is not None: + vllm_metrics_scraper.pause() eval_generate_time += time.monotonic() - gen_start validate_generator_output(len(generator_input["prompts"]), generator_output) generator_outputs.append(generator_output) @@ -142,6 +153,7 @@ async def evaluate_step_wise( cfg: SkyRLTrainConfig, global_step: int | None, tokenizer: AutoTokenizer, + vllm_metrics_scraper: Optional["VLLMMetricsScraper"] = None, ) -> Dict[str, float]: """Runs generation and evaluation of trajectories for step-wise training. @@ -154,6 +166,9 @@ async def evaluate_step_wise( global_step (int | None): current global step, or `None` to indicate a non-training context (e.g., eval-only) tokenizer (AutoTokenizer): tokenizer to use + vllm_metrics_scraper: when set, the open ``vllm/eval`` window is resumed + around each generation and paused after, so only generation time + counts toward eval throughput. Returns: Dict[str, float]: evaluation metrics @@ -178,7 +193,11 @@ async def evaluate_step_wise( global_step, ) gen_start = time.monotonic() + if vllm_metrics_scraper is not None: + vllm_metrics_scraper.resume() generator_output: GeneratorOutput = await generator.generate(generator_input) + if vllm_metrics_scraper is not None: + vllm_metrics_scraper.pause() eval_generate_time += time.monotonic() - gen_start traj_id_to_input = { traj_id.instance_id: {"env_class": env_class, "env_extras": env_extra} diff --git a/skyrl/train/trainer.py b/skyrl/train/trainer.py index 38d592e6e0..621a32c2bd 100644 --- a/skyrl/train/trainer.py +++ b/skyrl/train/trainer.py @@ -204,13 +204,18 @@ def _build_train_dataloader_and_compute_training_steps(self): self.total_training_steps = min(self.total_training_steps, self.cfg.trainer.max_training_steps) @torch.no_grad() - async def eval(self) -> Dict[str, float]: + async def eval(self, vllm_metrics_scraper: Optional[VLLMMetricsScraper] = None) -> Dict[str, float]: """ Run generation and scoring on the evaluation dataset. The eval metrics are recorded after having finished training `self.global_step` steps. Metrics recorded in global_step 0 corresponds to evaluations before training. + Args: + vllm_metrics_scraper: when provided, the eval loop calls + ``resume()``/``pause()`` around each generation so the scraper + attributes only generation time to the open ``vllm/eval`` window. + Returns: A dictionary of evaluation metrics. """ @@ -221,6 +226,7 @@ async def eval(self) -> Dict[str, float]: cfg=self.cfg, global_step=self.global_step, tokenizer=self.tokenizer, + vllm_metrics_scraper=vllm_metrics_scraper, ) else: eval_metrics = await evaluate( @@ -229,6 +235,7 @@ async def eval(self) -> Dict[str, float]: cfg=self.cfg, global_step=self.global_step, tokenizer=self.tokenizer, + vllm_metrics_scraper=vllm_metrics_scraper, ) return eval_metrics @@ -294,6 +301,13 @@ async def train(self): if not step_started: self._fire("on_step_start") step_started = True + # Open the train-rollout metrics window once per logical + # step; paused so only the generation spans count toward the + # throughput denominator (dynamic sampling may generate more + # than once before the step completes). + if self._vllm_metrics_scraper is not None: + await self._vllm_metrics_scraper.start("vllm/train") + self._vllm_metrics_scraper.pause() with Timer("step", self.all_timings): # for colocate_all=true, inference engine is always on GPU when starting the training step @@ -311,8 +325,12 @@ async def train(self): ) # 1.1. generation phase + if self._vllm_metrics_scraper is not None: + self._vllm_metrics_scraper.resume() with Timer("generate", self.all_timings): generator_output: GeneratorOutput = await self.generate(generator_input) + if self._vllm_metrics_scraper is not None: + self._vllm_metrics_scraper.pause() if self.cfg.generator.step_wise_trajectories: # NOTE: We use instance_ids from `trajectory_ids` here instead of re-using `uids` @@ -331,6 +349,13 @@ async def train(self): # if we are not continuing sampling, we sleep the inference engine await self.inference_engine_client.sleep() + # The train rollout for this step is done generating; close + # its metrics window. ``vllm/eval/*`` is collected separately + # around eval below. + vllm_metrics: Dict[str, float] = {} + if self._vllm_metrics_scraper is not None: + vllm_metrics = await self._vllm_metrics_scraper.stop() + # 1.2 postprocess rewards (and merge step-wise turns if enabled) with Timer("postprocess_generator_output", self.all_timings): generator_output, uids = self.postprocess_generator_output(generator_output, uids) @@ -434,27 +459,26 @@ async def train(self): or self.global_step == self.total_training_steps ) if force_eval or interval_eval: + # Open the eval-rollout window; the scraper itself measures + # the generation spans via resume()/pause() inside eval(). if self._vllm_metrics_scraper is not None: - await self._vllm_metrics_scraper.mark_pre_eval() + await self._vllm_metrics_scraper.start("vllm/eval") + self._vllm_metrics_scraper.pause() self._fire("on_eval_start") with Timer("eval", self.all_timings): - eval_metrics = await self.eval() + eval_metrics = await self.eval(vllm_metrics_scraper=self._vllm_metrics_scraper) self.all_metrics.update(eval_metrics) self._fire("on_eval_end", metrics=eval_metrics) + if self._vllm_metrics_scraper is not None: + vllm_metrics.update(await self._vllm_metrics_scraper.stop()) log_payload = { **self.all_metrics, **{f"timing/{k}": v for k, v in self.all_timings.items()}, + # vllm/train/* = train rollout, vllm/eval/* = eval rollout, + # each over its own generation time (owned by the scraper). + **vllm_metrics, } - if self._vllm_metrics_scraper is not None: - # vllm/* = train rollout / generate time; on eval steps - # vllm/eval/* = eval rollout / eval rollout time. - log_payload.update( - await self._vllm_metrics_scraper.sample_split( - generate_time_s=self.all_timings.get("generate"), - eval_generate_time_s=self.all_metrics.get("timing/eval_generate"), - ) - ) if self._ray_gpu_monitor is not None: log_payload.update(self._ray_gpu_monitor.flush()) diff --git a/skyrl/train/utils/vllm_metrics_scraper.py b/skyrl/train/utils/vllm_metrics_scraper.py index bcd4df17ba..e03f00facc 100644 --- a/skyrl/train/utils/vllm_metrics_scraper.py +++ b/skyrl/train/utils/vllm_metrics_scraper.py @@ -140,10 +140,21 @@ def discover_ray_metrics_urls() -> List[str]: class VLLMMetricsScraper: """Per-step snapshot of selected vLLM metrics from Ray's metrics agents. - ``sample()`` reports deltas vs. the previous step (used by the fully-async - trainer). ``mark_pre_eval()`` + ``sample_split()`` let the sync trainer - report the train rollout under ``vllm/*`` and the eval rollout under - ``vllm/eval/*`` instead of blending them. + Two ways to derive a window: + + * ``sample()`` reports deltas vs. the previous call against a wall-clock (or + caller-supplied generation) interval — used by the fully-async trainer. + * ``start(label)`` / ``pause()`` / ``resume()`` / ``stop()`` measure an + explicit window and own its timing. The scraper accumulates only the + un-paused time between ``start`` and ``stop``, so the caller never has to + thread a generation-time value back in or mark boundaries by ordering. + ``stop()`` returns the metrics nested under ``{label}/`` (e.g. + ``vllm/train/generation_throughput_tok_s``). The sync trainer uses this to + report the train rollout under ``vllm/train/*`` and the eval rollout under + ``vllm/eval/*`` instead of blending them. + + The two paths keep independent state, so a process may use either (the sync + trainer uses windows; the fully-async trainer uses ``sample()``). """ def __init__( @@ -155,9 +166,15 @@ def __init__( self._timeout = request_timeout_s self._prev_aggregated: Optional[Dict[str, float]] = None self._prev_timestamp: Optional[float] = None - self._mid_snapshot: Optional[Dict[str, float]] = None self._client: Optional[httpx.AsyncClient] = None self._warned_empty = False + # Explicit-window state (start/pause/resume/stop). ``_label is None`` + # means no window is open. + self._label: Optional[str] = None + self._window_prev: Optional[Dict[str, float]] = None + self._window_time_s: float = 0.0 + self._active_since: Optional[float] = None # start of the current un-paused span + self._paused: bool = False if not self._urls: logger.warning( "VLLMMetricsScraper: ray.nodes() returned no metrics endpoints; " @@ -240,41 +257,61 @@ async def sample(self, generation_time_s: Optional[float] = None) -> Dict[str, f self._prev_timestamp = now return out - async def mark_pre_eval(self) -> None: - """Snapshot the counters right before the eval rollout for :meth:`sample_split`.""" - self._mid_snapshot = await self._read_snapshot() + async def start(self, label: str) -> None: + """Open a metrics window labelled ``label`` (e.g. ``"vllm/train"``). - async def sample_split( - self, - *, - generate_time_s: Optional[float] = None, - eval_generate_time_s: Optional[float] = None, - ) -> Dict[str, float]: - """Sync-trainer sampling that separates train and eval rollouts. - - ``vllm/*`` covers the train rollout (previous step -> pre-eval mark); - when :meth:`mark_pre_eval` was called, ``vllm/eval/*`` covers the eval - rollout (mark -> now). Each is divided by its own generation time. + Snapshots the counters and starts the active-time clock. The window is + un-paused, so time accumulates immediately; call :meth:`pause` right + after if the work between ``start`` and the first generation should be + excluded from the throughput denominator. """ - cur = await self._read_snapshot() - mid = self._mid_snapshot - self._mid_snapshot = None - if cur is None: + if self._label is not None: + raise ValueError(f"`start({label!r})` called while window {self._label!r} is still open") + self._window_prev = await self._read_snapshot() + self._label = label + self._window_time_s = 0.0 + self._active_since = time.monotonic() + self._paused = False + + def pause(self) -> None: + """Stop accumulating active time until the next :meth:`resume`.""" + if self._label is None: + raise ValueError("`pause` called without an open window") + if self._paused: + raise ValueError("`pause` called without `resume`") + self._window_time_s += time.monotonic() - self._active_since + self._active_since = None + self._paused = True + + def resume(self) -> None: + """Resume accumulating active time after a :meth:`pause`.""" + if self._label is None: + raise ValueError("`resume` called without an open window") + if not self._paused: + raise ValueError("`resume` called without `pause`") + self._active_since = time.monotonic() + self._paused = False + + async def stop(self) -> Dict[str, float]: + """Close the window and return its metrics nested under ``{label}/``. + + The throughput denominator is the active (un-paused) time accumulated + between :meth:`start` and now. Returns ``{}`` when no endpoints are + configured. + """ + if self._label is None: + raise ValueError("`stop` called without an open window") + new_snapshot = await self._read_snapshot() + if not self._paused: + self._window_time_s += time.monotonic() - self._active_since + label, prev, window = self._label, self._window_prev, self._window_time_s + self._label = None + self._window_prev = None + self._active_since = None + self._paused = False + if new_snapshot is None: return {} - - # Train rollout ends at the pre-eval mark when eval ran, else at cur. - train_end = mid if mid is not None else cur - train_window = generate_time_s if (generate_time_s is not None and generate_time_s > 0) else None - out = self._window_metrics(self._prev_aggregated, train_end, train_window, "vllm/") - - if mid is not None: - eval_window = ( - eval_generate_time_s if (eval_generate_time_s is not None and eval_generate_time_s > 0) else None - ) - out.update(self._window_metrics(mid, cur, eval_window, "vllm/eval/")) - - self._prev_aggregated = cur - return out + return self._window_metrics(prev, new_snapshot, window, f"{label}/") @classmethod def _window_metrics( diff --git a/tests/train/test_vllm_metrics_scraper.py b/tests/train/test_vllm_metrics_scraper.py index 1307f1e6b1..59703f9576 100644 --- a/tests/train/test_vllm_metrics_scraper.py +++ b/tests/train/test_vllm_metrics_scraper.py @@ -372,86 +372,142 @@ def test_scraper_with_no_urls_is_noop(): assert out == {} -@pytest.mark.asyncio -async def test_sample_split_separates_train_and_eval_rollouts(): - """vllm/* uses the train window; vllm/eval/* uses the eval window. +def _split_snap(gen, prompt): + return _snapshot( + running=1, + waiting=0, + kv=0.4, + prefix_q=0, + prefix_h=0, + prompt_toks=prompt, + gen_toks=gen, + ttft_sum=0, + ttft_count=0, + itl_sum=0, + itl_count=0, + ) + - baseline gen=500 prompt=1000 - pre-eval gen=600 prompt=1100 (train rollout: +100 gen, +100 prompt) - log-time gen=1100 prompt=1300 (eval rollout: +500 gen, +200 prompt) - with generate_time=2s and eval_generate_time=5s. +@pytest.mark.asyncio +async def test_window_separates_train_and_eval_rollouts(): + """Separate start/stop windows label train and eval rollouts independently. + + train start gen=500 prompt=1000 + train stop gen=600 prompt=1100 (train rollout: +100 gen, +100 prompt) + eval start gen=600 prompt=1100 + eval stop gen=1100 prompt=1300 (eval rollout: +500 gen, +200 prompt) + The scraper owns timing: 2s of active train time, 5s of active eval time + (with prep/scoring paused out). """ scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) - def snap(gen, prompt): - return _snapshot( - running=1, - waiting=0, - kv=0.4, - prefix_q=0, - prefix_h=0, - prompt_toks=prompt, - gen_toks=gen, - ttft_sum=0, - ttft_count=0, - itl_sum=0, - itl_count=0, - ) - - texts = iter([snap(500, 1000), snap(600, 1100), snap(1100, 1300)]) + # One fetch per start() and per stop(). + texts = iter([_split_snap(500, 1000), _split_snap(600, 1100), _split_snap(600, 1100), _split_snap(1100, 1300)]) async def fake_fetch_all(): return parse_metrics_text(next(texts)) - with patch.object(scraper, "_fetch_all", fake_fetch_all): - # Step before eval: establishes the baseline (no throughput yet). - first = await scraper.sample_split(generate_time_s=1.0) - assert "vllm/generation_throughput_tok_s" not in first + # monotonic() is read at: train start, train stop, eval start, eval pause, + # eval resume, eval pause. Active train time = 102-100 = 2s; active eval + # time = (215-210) = 5s (the 200->200 prep span is paused out). + times = iter([100.0, 102.0, 200.0, 200.0, 210.0, 215.0]) + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + # Train rollout: a single generation spans the whole window. + await scraper.start("vllm/train") + train = await scraper.stop() - # Eval step: mark before eval, then sample after. - await scraper.mark_pre_eval() - out = await scraper.sample_split(generate_time_s=2.0, eval_generate_time_s=5.0) + # Eval rollout: paused for prep, resumed only around generation. + await scraper.start("vllm/eval") + scraper.pause() + scraper.resume() + scraper.pause() + eval_out = await scraper.stop() # Train rollout: (600-500)/2 == 50, (1100-1000)/2 == 50. - assert out["vllm/generation_throughput_tok_s"] == pytest.approx(50.0) - assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(50.0) + assert train["vllm/train/generation_throughput_tok_s"] == pytest.approx(50.0) + assert train["vllm/train/prompt_throughput_tok_s"] == pytest.approx(50.0) + assert not any(k.startswith("vllm/eval/") for k in train) # Eval rollout: (1100-600)/5 == 100, (1300-1100)/5 == 40. - assert out["vllm/eval/generation_throughput_tok_s"] == pytest.approx(100.0) - assert out["vllm/eval/prompt_throughput_tok_s"] == pytest.approx(40.0) + assert eval_out["vllm/eval/generation_throughput_tok_s"] == pytest.approx(100.0) + assert eval_out["vllm/eval/prompt_throughput_tok_s"] == pytest.approx(40.0) @pytest.mark.asyncio -async def test_sample_split_without_eval_emits_only_train_metrics(): - """On a non-eval step (no mark_pre_eval), no vllm/eval/* keys appear.""" - scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) +async def test_window_accumulates_active_time_across_multiple_generations(): + """pause/resume around each generation sums only the active spans. - def snap(gen, prompt): - return _snapshot( - running=1, - waiting=0, - kv=0.4, - prefix_q=0, - prefix_h=0, - prompt_toks=prompt, - gen_toks=gen, - ttft_sum=0, - ttft_count=0, - itl_sum=0, - itl_count=0, - ) + Mirrors a dynamic-sampling step (or an eval loop) that generates more than + once: the paused gap between generations is excluded from the denominator. + """ + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) - texts = iter([snap(0, 0), snap(200, 100)]) + texts = iter([_split_snap(0, 0), _split_snap(300, 100)]) async def fake_fetch_all(): return parse_metrics_text(next(texts)) + # monotonic() reads: start(0), pause(0), resume(10), pause(12) -> gen1 = 2s, + # resume(112), pause(114) -> gen2 = 2s. The 100s paused gap (12 -> 112) is + # excluded, so active time is 4s, not 114s. + times = iter([0.0, 0.0, 10.0, 12.0, 112.0, 114.0]) + + with ( + patch.object(scraper, "_fetch_all", fake_fetch_all), + patch( + "skyrl.train.utils.vllm_metrics_scraper.time.monotonic", + side_effect=lambda: next(times), + ), + ): + await scraper.start("vllm/eval") + scraper.pause() + scraper.resume() # generation 1 + scraper.pause() + scraper.resume() # generation 2 + scraper.pause() + out = await scraper.stop() + + # gen 300 over 4s active == 75 tok/s; prompt 100 over 4s == 25 tok/s. + assert out["vllm/eval/generation_throughput_tok_s"] == pytest.approx(75.0) + assert out["vllm/eval/prompt_throughput_tok_s"] == pytest.approx(25.0) + + +@pytest.mark.asyncio +async def test_window_pause_resume_misuse_raises(): + scraper = VLLMMetricsScraper(urls=["http://stub/metrics"]) + + async def fake_fetch_all(): + return parse_metrics_text(_split_snap(0, 0)) + with patch.object(scraper, "_fetch_all", fake_fetch_all): - await scraper.sample_split(generate_time_s=1.0) - out = await scraper.sample_split(generate_time_s=4.0) + with pytest.raises(ValueError): + scraper.pause() # no open window + with pytest.raises(ValueError): + await scraper.stop() # no open window + await scraper.start("vllm/train") + with pytest.raises(ValueError): + await scraper.start("vllm/eval") # window already open + scraper.pause() + with pytest.raises(ValueError): + scraper.pause() # double pause + scraper.resume() + with pytest.raises(ValueError): + scraper.resume() # double resume + await scraper.stop() + - assert out["vllm/generation_throughput_tok_s"] == pytest.approx(200.0 / 4.0) - assert out["vllm/prompt_throughput_tok_s"] == pytest.approx(100.0 / 4.0) - assert not any(k.startswith("vllm/eval/") for k in out) +@pytest.mark.asyncio +async def test_window_returns_empty_without_endpoints(): + scraper = VLLMMetricsScraper(urls=[]) + await scraper.start("vllm/train") + out = await scraper.stop() + assert out == {} @pytest.mark.asyncio From f7502157b4f2e8b8197fe534139a6c1b839e8e52 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 24 Jun 2026 05:02:12 +0000 Subject: [PATCH 23/28] [mtp] Fix C-full hidden() to actually exclude the MTP head from the policy grad-norm/clip MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Megatron's get_main_grads_for_grad_norm() collects a param's grad when `grad is not None` — it does NOT check requires_grad. The previous hidden() only flipped requires_grad, so the head's already- populated grad still entered the policy grad-norm, inflating policy/grad_norm (observed ~2.5 vs ~0.3 no-MTP) and over-clipping the policy via the shared max_grad_norm — a head->policy coupling through the clip. hidden() now stashes+clears grad/main_grad/decoupled_grad for the policy step and restores them before the separate MTP optimizer steps. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../skyrl_train/mtp/separate_optim.py | 33 ++++++++++++++----- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/skyrl/backends/skyrl_train/mtp/separate_optim.py b/skyrl/backends/skyrl_train/mtp/separate_optim.py index c1c45a3ef7..971b8db65a 100644 --- a/skyrl/backends/skyrl_train/mtp/separate_optim.py +++ b/skyrl/backends/skyrl_train/mtp/separate_optim.py @@ -183,21 +183,36 @@ def zero_grad_buffer(self) -> None: @contextlib.contextmanager def hidden(self): - """Temporarily set the head's params ``requires_grad=False`` so the POLICY optimizer's - grad-norm + clip EXCLUDE the head. Megatron computes the policy grad-norm by iterating the - GPTModel's ``requires_grad`` params and reading ``main_grad`` — and the head is *structurally* - still inside the GPTModel (so capture/replay/export work), so without this it would pick up the - head's grad (from its separate buffer) and dilute the policy clip exactly like the un-isolated - run. This is a pure read-time filter while the policy step runs; the head's grads stay in its - own buffer and are clipped/stepped separately by :meth:`step`.""" - saved = [(p, p.requires_grad) for p in self.mtp_params] + """Exclude the head from the POLICY optimizer's grad-norm + clip for the duration of the + policy step, then restore so the separate MTP optimizer (stepped right after) reduces/steps + the real grads. + + IMPORTANT: Megatron's ``get_main_grads_for_grad_norm()`` collects a param's grad whenever + ``grad is not None`` — it does NOT check ``requires_grad``. So flipping ``requires_grad`` alone + (the original implementation) does NOT remove the head from the policy grad-norm/clip: the head's + grad is already populated by backward, so it still gets counted and inflates ``policy/grad_norm``, + over-clipping the policy (a head→policy coupling through the shared clip). We therefore stash and + clear the grad-bearing attributes Megatron may read (``grad`` / ``main_grad`` / ``decoupled_grad``) + so the head is genuinely invisible to the policy norm, and restore them on exit (the underlying + grad-buffer data is untouched — we only detach the attribute references). ``requires_grad`` is + also cleared for any path that does honor it.""" + _GRAD_ATTRS = ("grad", "main_grad", "decoupled_grad") + saved = [] for p in self.mtp_params: + grads = {a: getattr(p, a, None) for a in _GRAD_ATTRS} + saved.append((p, p.requires_grad, grads)) p.requires_grad = False + for a in _GRAD_ATTRS: + if getattr(p, a, None) is not None: + setattr(p, a, None) try: yield finally: - for p, rg in saved: + for p, rg, grads in saved: p.requires_grad = rg + for a, g in grads.items(): + if g is not None: + setattr(p, a, g) def finalize_grads(self) -> None: """Reduce the head's accumulated grads: DP reduce-scatter (own buffer) + TP all-reduce of the From 30a4997528411bcaf8170119f27d18176e95ff70 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 24 Jun 2026 08:08:44 +0000 Subject: [PATCH 24/28] [mtp] Disable Megatron's native in-forward MTP loss via call-site monkey-patch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the MiMo entropy collapse / inflated policy grad-norm: GPTModel/HybridModel.forward call process_mtp_loss unconditionally when MTP heads are built, and a megatron-core update made it derive labels from input_ids ("e.g. RL training") — so SkyRL's "pass no labels to short-circuit it" no longer works. The native hard-CE MTP loss then back-props into the policy trunk (grad-norm 2.5 vs 0.3, entropy collapse), independent of mtp_loss_weight. Replace process_mtp_loss with a no-op at its call sites (gpt_model, hybrid_model) so only SkyRL's decoupled soft-CE loss trains the head. Applied at config time, before any forward. The patch raises loudly if Megatron renames/removes the function rather than silently re-enabling the native loss. Confirmed: with the native loss off, policy grad_norm drops 2.5 -> ~0.3 (matches no-MTP) and entropy stops collapsing, while the decoupled head loss keeps training. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../skyrl_train/mtp/native_loss_patch.py | 91 +++++++++++++++++++ .../megatron/megatron_model_wrapper.py | 7 +- .../workers/megatron/megatron_worker.py | 15 ++- 3 files changed, 109 insertions(+), 4 deletions(-) create mode 100644 skyrl/backends/skyrl_train/mtp/native_loss_patch.py diff --git a/skyrl/backends/skyrl_train/mtp/native_loss_patch.py b/skyrl/backends/skyrl_train/mtp/native_loss_patch.py new file mode 100644 index 0000000000..57024f9c6b --- /dev/null +++ b/skyrl/backends/skyrl_train/mtp/native_loss_patch.py @@ -0,0 +1,91 @@ +"""Disable Megatron's native in-forward MTP loss (``process_mtp_loss``). + +SkyRL trains the MTP/draft head with its OWN decoupled soft-CE distillation loss (see +``mtp/soft_ce.py`` and ``MegatronModelWrapper``), NOT Megatron's pretraining-style MTP auxiliary +loss. But ``GPTModel.forward`` / ``HybridModel.forward`` call ``process_mtp_loss`` unconditionally +whenever the model is built with MTP heads (``mtp_num_layers`` set) and run in training/eval. That +native loss is a hard cross-entropy whose gradient — unless ``mtp_detach_heads`` is set — flows +straight into the shared trunk and output embedding, corrupting the RL policy (inflated grad-norm +and entropy collapse), with a magnitude set by Megatron's own scaling (so it is independent of +SkyRL's ``mtp_loss_weight``). + +Megatron only short-circuits the native loss when BOTH ``labels`` and ``input_ids`` are None. A +megatron-core update added "derive labels from ``input_ids`` (e.g. RL training)", so passing no +labels no longer disables it. Rather than depend on that label / forward-gating behaviour (which is +exactly what silently broke), we replace ``process_mtp_loss`` at its call sites with a no-op that +returns the main-model hidden chunk and computes no loss. SkyRL's decoupled head training is +unaffected — the head still trains via SkyRL's explicit loss. + +Robustness: this patch is loud, not silent. It raises if a target module no longer exposes +``process_mtp_loss`` (a Megatron rename). Its one blind spot is Megatron *inlining* the loss into +the forward instead of calling the function — which the grad-isolation acceptance test (the policy +gradient must match a no-MTP build, see ``tests/.../megatron/test_mtp_grad_coupling.py``) is there +to catch. +""" + +from __future__ import annotations + +import importlib + +import torch + +# Modules whose ``forward`` calls ``process_mtp_loss``. Each does ``from ...multi_token_prediction +# import process_mtp_loss``, binding the name in its OWN namespace, so we must patch each call +# site's module (patching the definition module would not rebind the already-imported names). +_PATCH_TARGETS = ( + "megatron.core.models.gpt.gpt_model", + "megatron.core.models.hybrid.hybrid_model", +) +_ATTR = "process_mtp_loss" +_SENTINEL = "_skyrl_native_mtp_loss_disabled" + + +def _skyrl_skip_native_mtp_loss(hidden_states, *args, **kwargs): + """No-op replacement for Megatron's ``process_mtp_loss``. + + Megatron splits ``hidden_states`` into ``1 + mtp_num_layers`` chunks along dim 0 and returns the + first (the main-model hidden states) after applying the MTP loss via ``MTPLossAutoScaler``. We + return that same first chunk WITHOUT computing or applying any loss, so no native-MTP gradient + reaches the model. Both known call sites pass ``hidden_states`` and ``config`` as kwargs; the + positional fallback is defensive only. + """ + if hidden_states is None: + hidden_states = kwargs.get("hidden_states") + config = kwargs.get("config") + if config is None: + config = next((a for a in args if hasattr(a, "mtp_num_layers")), None) + num_layers = getattr(config, "mtp_num_layers", None) if config is not None else None + if not num_layers: + return hidden_states + return torch.chunk(hidden_states, 1 + num_layers, dim=0)[0] + + +def disable_native_mtp_loss() -> None: + """Replace ``process_mtp_loss`` at every known call site with a no-op (idempotent). + + Raises ``RuntimeError`` if a target module exists but no longer exposes ``process_mtp_loss`` + (Megatron renamed/removed it) or if no target module can be imported at all — so the breakage + is loud instead of silently re-enabling the native loss. + """ + patched_any = False + for mod_name in _PATCH_TARGETS: + try: + mod = importlib.import_module(mod_name) + except ImportError: + continue # not present in this Megatron build (e.g. no hybrid models) + if getattr(mod, _SENTINEL, False): + patched_any = True + continue + if not hasattr(mod, _ATTR): + raise RuntimeError( + f"Cannot disable native MTP loss: '{mod_name}' no longer exposes '{_ATTR}'. " + "Megatron's MTP-loss API changed — update mtp/native_loss_patch.py." + ) + setattr(mod, _ATTR, _skyrl_skip_native_mtp_loss) + setattr(mod, _SENTINEL, True) + patched_any = True + if not patched_any: + raise RuntimeError( + f"Cannot disable native MTP loss: none of {_PATCH_TARGETS} could be imported. " + "Megatron's MTP module layout changed — update mtp/native_loss_patch.py." + ) diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index a239ee31eb..3a56c0bab1 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -348,9 +348,10 @@ def forward_backward_mini_batch( # Multi-Token Prediction (MTP): if the model was built with native MTP heads, train them with # an explicit decoupled loss instead of Megatron's in-forward process_mtp_loss path. The heads - # still run inside the forward (so we reuse their rotary embeddings) but with NO labels, so - # process_mtp_loss short-circuits and no MTP gradient couples onto the trunk; a forward hook - # captures their hidden states (trunk input optionally detached) for us to score. Training only. + # still run inside the forward (so we reuse their rotary embeddings); the native process_mtp_loss + # is disabled at its call sites (see mtp/native_loss_patch.py, applied at config time) so no + # native MTP gradient couples onto the trunk. A forward hook captures the heads' hidden states + # (trunk input optionally detached) for us to score. Training only. model_config = get_model_config(self.actor_module[0]) mtp_enabled = (not forward_only) and bool(getattr(model_config, "mtp_num_layers", None)) mcfg = self.cfg.policy.megatron_config diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index c011e37186..4edb6a6ce3 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -477,11 +477,24 @@ def init_configs( # loss (see MegatronModelWrapper) instead of Megatron's in-forward # process_mtp_loss / MTPLossAutoScaler path. The explicit weight is # megatron_config.mtp_loss_weight. + # + # CRITICAL: GPTModel/HybridModel.forward call process_mtp_loss unconditionally when MTP + # heads exist, and Megatron now DERIVES labels from input_ids ("e.g. RL training") so + # passing no labels no longer short-circuits it. Left active, that native hard-CE loss + # back-props into the policy trunk (inflated grad-norm, entropy collapse), independent of + # mtp_loss_weight. Patch process_mtp_loss to a no-op at its call sites so only SkyRL's + # decoupled loss trains the head. Must run before any forward (here, at config time). + from skyrl.backends.skyrl_train.mtp.native_loss_patch import ( + disable_native_mtp_loss, + ) + + disable_native_mtp_loss() logger.info( f"MTP enabled (decoupled): mtp_num_layers={provider.mtp_num_layers}, " f"mtp_loss_type={megatron_config.mtp_loss_type}, " f"mtp_loss_weight={megatron_config.mtp_loss_weight}, " - f"mtp_detach_trunk={megatron_config.mtp_detach_trunk}" + f"mtp_detach_trunk={megatron_config.mtp_detach_trunk} " + "(native process_mtp_loss disabled)" ) provider.finalize() From eae60f546f5f3b302cb3d909690aa3fca0a01c78 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Wed, 24 Jun 2026 16:36:46 +0000 Subject: [PATCH 25/28] [mtp] Add SKYRL_DISABLE_SPEC debug guard (train MTP head, plain rollout) Co-Authored-By: Claude Opus 4.8 (1M context) --- skyrl/train/utils/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index f63207a1d3..04978f6f37 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -254,6 +254,12 @@ def _apply_mtp_config(cfg: SkyRLTrainConfig): mcfg.mtp_loss_type = mtp.loss_type mcfg.mtp_loss_weight = mtp.loss_weight + # DEBUG GUARD: SKYRL_DISABLE_SPEC=1 keeps MTP head TRAINING on (the decoupled draft loss set + # above) but leaves vLLM rollout in plain autoregressive (lossless) decode — isolates the + # MTP-training machinery from the spec-decode rollout sampler. + if os.environ.get("SKYRL_DISABLE_SPEC") == "1": + return + # Inference side: vLLM MTP speculative decoding with the same draft depth. Don't clobber an # explicit user-provided speculative_config. ie_cfg = cfg.generator.inference_engine From d1b89d91e7b6e9394b2310ef90932a77f17cd29d Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Thu, 25 Jun 2026 06:02:38 +0000 Subject: [PATCH 26/28] adjust training scripts to dedicated folder --- ...run_megatron_dapo_mimo_7b_rl_specdecode.sh | 0 ...run_megatron_dapo_qwen3.5_2b_specdecode.sh | 0 ...run_megatron_dapo_qwen3.5_9b_specdecode.sh | 0 ...tron_dapo_qwen3.5_9b_specdecode_notrain.sh | 0 .../run_search_megatron_mimo_7b_specdecode.sh | 160 ++++++++++++++++++ 5 files changed, 160 insertions(+) rename examples/train/{megatron => spec_decode}/run_megatron_dapo_mimo_7b_rl_specdecode.sh (100%) rename examples/train/{megatron => spec_decode}/run_megatron_dapo_qwen3.5_2b_specdecode.sh (100%) rename examples/train/{megatron => spec_decode}/run_megatron_dapo_qwen3.5_9b_specdecode.sh (100%) rename examples/train/{megatron => spec_decode}/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh (100%) create mode 100644 examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh diff --git a/examples/train/megatron/run_megatron_dapo_mimo_7b_rl_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh similarity index 100% rename from examples/train/megatron/run_megatron_dapo_mimo_7b_rl_specdecode.sh rename to examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh similarity index 100% rename from examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh rename to examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh similarity index 100% rename from examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh rename to examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh similarity index 100% rename from examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh rename to examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh diff --git a/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh b/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh new file mode 100644 index 0000000000..3719ab7273 --- /dev/null +++ b/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh @@ -0,0 +1,160 @@ +set -x + +# Colocated GRPO multi-turn SearchR1 training+generation for XiaomiMiMo/MiMo-7B-RL (dense) with +# Megatron, WITH native Multi-Token Prediction (MTP) speculative decoding (k=3) for faster rollout. +# +# This is the MiMo-7B-RL + Megatron + spec-decode counterpart of examples/train/search/run_search.sh +# (which trains Qwen2.5-3B-Instruct with FSDP and NO spec decode). The SearchR1 recipe knobs +# (dataset, search env, multi-turn sampling, search server) are kept identical to run_search.sh / +# run_search_megatron.sh; only the model, the FSDP->Megatron strategy, and the MTP block differ. +# +# MTP on MiMo-7B-RL (see run_megatron_dapo_mimo_7b_rl_specdecode.sh for the full rationale): +# - MiMo ships 1 native MTP head (`num_nextn_predict_layers: 1`), so both the training-side head +# and the inference-side drafter come from the same pretrained weights (no cold-start head). +# - Plain Qwen2-style dense backbone (no GDN), so megatron sample packing IS supported +# (remove_microbatch_padding=true) and no special vLLM prefill backend is needed. +# - Untied embeddings => the decoupled draft loss (mtp_detach_shared_output default true) trains +# ONLY the MTP-head params; the draft gradient never nudges the policy backbone. +# - vLLM MTP speculative decoding draft depth = k = num_speculative_tokens. Here k=3: the single +# trained head is reused autoregressively at draft time. Acceptance is logged as +# `vllm/draft_acceptance_rate`. +# +# Runs on 1 node of 8xH100s (80GB each). +# +# Setup (see examples/train/search/README.md): +# 1. Dataset parquet: +# uv run --isolated examples/train/search/searchr1_dataset.py --local_dir /mnt/local_storage/data/searchR1 +# 2. Download the wiki-18 e5 index + corpus into the same dir, assemble e5_Flat.index, gunzip corpus. +# 3. Start the local e5 retrieval server (separate faiss-gpu conda env) on :8000. +# Then launch: +# export WANDB_API_KEY= +# bash examples/train/megatron/run_search_megatron_mimo_7b_specdecode.sh + +MODEL_NAME="XiaomiMiMo/MiMo-7B-RL" +# Dataset on the fast, non-persistent local disk (not the ~/default quota). +DATA_DIR="/mnt/local_storage/data/searchR1" + +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 + +# megatron config -- MiMo-7B-RL is a dense model, so no expert parallelism. +# TP=4, PP=1, CP=1 => DP=2; TP stays within the single-node NVLink domain. +MEGATRON_TP=4 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=null + +# One vLLM engine per GPU (TP=1). A single MiMo-7B copy is ~14GB bf16, fits per engine alongside +# KV cache + the small MTP drafter at gpu_memory_utilization=0.5 (policy offloaded during gen under +# colocate_all). gmu=0.5 also leaves headroom for the faiss-gpu retriever server (~6GB/GPU). +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TP=1 + +MICRO_TRAIN_BATCH_SIZE_PER_GPU=1 +MICRO_FORWARD_BATCH_SIZE_PER_GPU=2 + +# TIS parameters (match run_search.sh) +TIS_TYPE=token +TIS_IMP_RATIO_CAP=2.0 + +# Multi-Token Prediction (MTP) speculative decoding -- k=3. +# trainer.mtp is the single high-level knob; validate_cfg propagates it to the training side +# (policy.megatron_config.mtp_*) and the inference side (speculative_config). mtp_num_layers is left +# at default => inferred from MiMo's HF config (1 native head). k>1 reuses that head autoregressively. +MTP_ENABLED=true +MTP_NUM_SPECULATIVE_TOKENS=3 +MTP_LOSS_TYPE="soft_ce" # distill against the policy's own detached next-token distribution +MTP_LOSS_WEIGHT=0.5 +MTP_LOSS_TOPK=256 # top-k draft loss: O(seq*k) memory vs O(seq*vocab) + +# MiMo flags -- plain Qwen2-style dense attention (no GDN): sample packing supported, no special +# vLLM prefill backend needed. +REMOVE_MICROBATCH_PADDING=true +DISTRIBUTED_EXECUTOR_BACKEND="mp" +ENFORCE_EAGER=true # cuda graphs can cause instability with weight sync +export _SKYRL_USE_NEW_INFERENCE=0 +export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 + +RUN_NAME="sd_search_mimo_7b_rl_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}_k${MTP_NUM_SPECULATIVE_TOKENS}" + +uv run --isolated --extra megatron -m skyrl.train.entrypoints.main_base \ + data.train_data="['${DATA_DIR}/train.parquet']" \ + data.val_data="['${DATA_DIR}/validation.parquet']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.policy.optimizer_config.lr=1.0e-6 \ + trainer.policy.optimizer_config.max_grad_norm=0.5 \ + trainer.policy.optimizer_config.num_warmup_steps=94 \ + trainer.algorithm.use_kl_loss=true \ + trainer.algorithm.kl_loss_coef=0.001 \ + trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ + trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.placement.colocate_all=true \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.ref_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TP \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ + generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ + generator.inference_engine.gpu_memory_utilization=0.5 \ + trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ + trainer.epochs=1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=512 \ + trainer.policy_mini_batch_size=256 \ + trainer.micro_forward_batch_size_per_gpu=$MICRO_FORWARD_BATCH_SIZE_PER_GPU \ + trainer.micro_train_batch_size_per_gpu=$MICRO_TRAIN_BATCH_SIZE_PER_GPU \ + trainer.max_prompt_length=2048 \ + generator.max_input_length=4096 \ + generator.sampling_params.max_generate_length=500 \ + generator.inference_engine.async_engine=true \ + generator.batched=false \ + generator.use_conversation_multi_turn=false \ + generator.n_samples_per_prompt=5 \ + generator.max_turns=4 \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.stop='["", ""]' \ + environment.env_class="search" \ + environment.skyrl_gym.max_env_workers=16 \ + environment.skyrl_gym.search.log_requests=false \ + environment.skyrl_gym.search.search_url="http://127.0.0.1:8000/retrieve" \ + environment.skyrl_gym.search.topk=3 \ + trainer.mtp.enabled=$MTP_ENABLED \ + trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ + trainer.mtp.loss_type=$MTP_LOSS_TYPE \ + trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ + trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ + trainer.logger="wandb" \ + trainer.project_name="mimo_7b_rl_searchr1" \ + trainer.run_name="${RUN_NAME}" \ + trainer.ckpt_interval=20 \ + trainer.hf_save_interval=100 \ + trainer.max_ckpts_to_keep=3 \ + trainer.resume_mode=latest \ + trainer.ckpt_path="/mnt/local_storage/ckpts/${RUN_NAME}" \ + trainer.eval_batch_size=256 \ + trainer.eval_before_train=false \ + trainer.eval_interval=50 \ + generator.eval_sampling_params.temperature=0 \ + generator.eval_sampling_params.stop='["", ""]' \ + generator.eval_sampling_params.max_generate_length=500 \ + trainer.export_path="/mnt/local_storage/exports/${RUN_NAME}" \ + $@ From edfd8eb7ebc2f45e22c44a812e8c049c5c915b69 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Thu, 25 Jun 2026 21:11:02 +0000 Subject: [PATCH 27/28] remove 2b script, update comments --- .../megatron/run_megatron_dapo_qwen3.5_2b.sh | 149 ------------ ...run_megatron_dapo_mimo_7b_rl_specdecode.sh | 9 +- ...run_megatron_dapo_qwen3.5_2b_specdecode.sh | 207 ----------------- ...tron_dapo_qwen3.5_9b_specdecode_notrain.sh | 213 ------------------ .../vllm/spec_decode_metrics.py | 35 ++- .../skyrl_train/inference_servers/utils.py | 7 +- .../inference_servers/vllm_worker.py | 3 +- .../skyrl_train/mtp/hidden_capture.py | 26 +-- .../mtp/{separate_optim.py => mtp_optim.py} | 64 ++---- .../skyrl_train/mtp/native_loss_patch.py | 27 +-- .../workers/megatron/megatron_worker.py | 14 +- skyrl/train/config/config.py | 2 +- .../mtp/test_build_vllm_cli_args_mtp.py | 1 + 13 files changed, 77 insertions(+), 680 deletions(-) delete mode 100644 examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh delete mode 100644 examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh delete mode 100644 examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh rename skyrl/backends/skyrl_train/mtp/{separate_optim.py => mtp_optim.py} (75%) diff --git a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh b/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh deleted file mode 100644 index e7a3128a91..0000000000 --- a/examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh +++ /dev/null @@ -1,149 +0,0 @@ -set -x - -# Colocated DAPO training+generation for Qwen3.5-2B (dense) on DAPO with Megatron. -# Runs on 1 node of 8xH100s (80GB each). -# -# NOTE: verify the exact HF repo id for the 2B model before running -# (e.g. `hf download Qwen/Qwen3.5-2B` / check https://huggingface.co/Qwen). -# -# Prepare data onto the fast local disk first: -# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh -# Then launch: -# bash examples/train/megatron/run_megatron_dapo_qwen3.5_2b.sh - -MODEL_NAME="Qwen/Qwen3.5-2B" -# Use the fast, non-persistent local disk for data (not the ~/default quota). -DATA_DIR="/mnt/local_storage/data/dapo" -TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" -TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" -NUM_NODES=1 -NUM_GPUS_PER_NODE=8 -# 2B is small: TP=1 inference per engine (no TP comm), one engine per GPU. -NUM_INFERENCE_ENGINES=8 -INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 -LOGGER="wandb" # change to "console" to print to stdout - -CLIP_RATIO_LOW=0.2 -CLIP_RATIO_HIGH=0.28 -# use token mean loss reduction -LOSS_REDUCTION="token_mean" -# applies overlong filtering (but not soft overlong punishment) -APPLY_OVERLONG_FILTERING=true -# apply soft overlong punishment with custom trainer impl in main_dapo.py -OVERLONG_BUFFER_LEN=$((1024 * 4)) -OVERLONG_BUFFER_PENALTY_FACTOR=1.0 - -# other DAPO parameters -USE_KL_LOSS=false -TEMPERATURE=1.0 -TOP_P=1.0 -EVAL_TOP_P=0.7 -CLIP_RATIO_C=10.0 -MAX_PROMPT_LENGTH=$((1024 * 2)) -MAX_RESPONSE_LENGTH=$((1024 * 8)) - -# repro run parameters -# Quality-comparison profile: larger batch lowers reward-curve variance (~4x vs batch 32) -# so a spec-decode quality effect is detectable. rollout = 128 * 8 = 1024 seqs. -# MINI_BATCH_SIZE == TRAIN_BATCH_SIZE => 1 on-policy update/rollout. Keep IDENTICAL across no-spec/spec runs. -TRAIN_BATCH_SIZE=128 -MINI_BATCH_SIZE=32 -N_SAMPLES_PER_PROMPT=8 -EVAL_N_SAMPLES_PER_PROMPT=16 -ENFORCE_EAGER=true # cuda graphs can cause some instability -LR=1e-6 - -# megatron config -- Qwen3.5-2B is a dense model, so no expert parallelism. -# TP=2 (not 1): with 8K-token responses, TP=2 auto-enables sequence parallelism, -# which shards activations/vocab-logits across the TP group. TP=1 keeps full -# activations on each GPU and OOMs in the backward grad-sync at gpu_mem_util=0.5. -# TP=2, PP=1, CP=1 => DP=4. -MEGATRON_TP=2 -MEGATRON_PP=1 -MEGATRON_CP=1 -MEGATRON_EP=1 -MEGATRON_ETP=null - - -# TIS parameters -TIS_IMP_RATIO_CAP=2.0 -TIS_TYPE=token - - -# Qwen3.5 flags -REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 -ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 -DISTRIBUTED_EXECUTOR_BACKEND="mp" -export _SKYRL_USE_NEW_INFERENCE=0 -export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 - -uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ - data.train_data="['$TRAIN_FILE']" \ - data.val_data="['$TEST_FILE']" \ - trainer.algorithm.advantage_estimator="grpo" \ - trainer.algorithm.policy_loss_type="dual_clip" \ - trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ - trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ - trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ - generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ - generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ - generator.sampling_params.temperature=$TEMPERATURE \ - generator.sampling_params.top_p=$TOP_P \ - generator.eval_sampling_params.top_p=$EVAL_TOP_P \ - generator.eval_sampling_params.temperature=$TEMPERATURE \ - generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ - trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ - trainer.policy.model.path="$MODEL_NAME" \ - trainer.placement.colocate_all=true \ - trainer.strategy=megatron \ - generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ - trainer.placement.policy_num_nodes=$NUM_NODES \ - trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ - generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ - generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ - generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ - trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ - trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ - trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ - trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ - trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ - trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ - trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ - trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ - trainer.epochs=10 \ - trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ - trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ - trainer.eval_batch_size=1024 \ - trainer.eval_before_train=false \ - trainer.eval_interval=5 \ - trainer.update_epochs_per_batch=1 \ - trainer.train_batch_size=$TRAIN_BATCH_SIZE \ - trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ - trainer.micro_forward_batch_size_per_gpu=1 \ - trainer.micro_train_batch_size_per_gpu=1 \ - trainer.ckpt_interval=-1 \ - trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ - generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.policy.optimizer_config.lr=$LR \ - trainer.policy.optimizer_config.num_warmup_steps=5 \ - trainer.policy.optimizer_config.weight_decay=0.1 \ - trainer.policy.optimizer_config.max_grad_norm=1.0 \ - generator.inference_engine.backend=vllm \ - generator.inference_engine.run_engines_locally=true \ - generator.inference_engine.weight_sync_backend=nccl \ - generator.inference_engine.async_engine=true \ - generator.batched=true \ - environment.env_class=aime \ - generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ - generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ - generator.inference_engine.gpu_memory_utilization=0.5 \ - trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo_sd" \ - trainer.run_name="nosd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.export_path="/mnt/local_storage/exports/dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.hf_save_interval=300 \ - trainer.resume_mode=latest \ - trainer.max_ckpts_to_keep=3 \ - trainer.ckpt_path="/mnt/local_storage/ckpts/dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - $@ diff --git a/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh index b0aeecf679..481d5e9deb 100644 --- a/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh +++ b/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh @@ -51,7 +51,7 @@ NUM_GPUS_PER_NODE=8 # (the policy is offloaded during generation under colocate_all). NUM_INFERENCE_ENGINES=8 INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 -LOGGER="wandb" # change to "console" to print to stdout +LOGGER="${LOGGER:-wandb}" # change to "console" to print to stdout (or set LOGGER env var) CLIP_RATIO_LOW=0.2 CLIP_RATIO_HIGH=0.28 @@ -113,7 +113,7 @@ TIS_TYPE=token MTP_ENABLED=true MTP_NUM_SPECULATIVE_TOKENS=3 MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) -MTP_LOSS_WEIGHT=0.5 +MTP_LOSS_WEIGHT=0.2 # NOTE: trainer.policy.megatron_config.mtp_detach_shared_output defaults to true. MiMo has UNTIED # embeddings, so this detaches the (separate) output_layer.weight passed into the draft projection, # along with the MTP block's re-embedding, so the draft gradient trains ONLY the MTP-head params and @@ -122,6 +122,10 @@ MTP_LOSS_WEIGHT=0.5 # Top-k draft loss: distill only the teacher's top-k tokens instead of the full vocab, keeping # draft-loss memory at O(seq*k) vs O(seq*vocab). Only applies to mtp_loss_type="soft_ce". Here k=256. MTP_LOSS_TOPK=256 +# Train the MTP head in the policy's native DDP buffer + optimizer (no separate C-full buffer); the +# native MTP loss stays patched off and the draft loss stays autograd-decoupled, so no MTP grad +# reaches the policy -- only the ~ulp-level (DP>=3) grad-buffer reduction reordering is given up. +MTP_SEPARATE_OPTIMIZER=false # MiMo flags -- plain Qwen2-style dense attention (no GDN), so sample packing is supported and no @@ -200,6 +204,7 @@ uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ trainer.mtp.loss_type=$MTP_LOSS_TYPE \ trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ + trainer.policy.megatron_config.mtp_separate_optimizer=$MTP_SEPARATE_OPTIMIZER \ trainer.logger="$LOGGER" \ trainer.project_name="mimo_7b_rl_dapo" \ trainer.run_name="sd_dapo_mimo_7b_rl_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ diff --git a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh deleted file mode 100644 index 80bb179000..0000000000 --- a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_2b_specdecode.sh +++ /dev/null @@ -1,207 +0,0 @@ -set -x - -# Colocated DAPO training+generation for Qwen3.5-2B (dense) on DAPO with Megatron, -# WITH Multi-Token Prediction (MTP) speculative decoding for faster rollout. -# -# This is the spec-decode counterpart of run_megatron_dapo_qwen3.5_2b.sh. Every -# knob below is IDENTICAL to the no-spec script (same batch sizes, LR, parallelism, -# sampling) so a reward-curve / throughput comparison is apples-to-apples. The ONLY -# difference is the `trainer.mtp.*` block at the bottom. -# -# What MTP on does (single high-level `trainer.mtp` knob, see skyrl/train/config/config.py): -# - Training side: builds + trains Qwen3.5-2B's native MTP head (the model ships -# `mtp_num_hidden_layers: 1`) with a decoupled draft loss (soft-CE distillation against -# the policy's own detached next-token distribution). The draft gradient never pulls on -# the policy backbone. -# - Inference side: enables vLLM MTP speculative decoding -# (`speculative_config={"method": "mtp", "num_speculative_tokens": 1}`). vLLM loads the -# MTP head from the same policy checkpoint, and SkyRL's weight sync keeps the draft head -# in sync with the trained policy each step. -# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. -# -# Runs on 1 node of 8xH100s (80GB each). -# -# Prepare data onto the fast local disk first: -# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh -# Then launch: -# bash examples/train/megatron/run_megatron_dapo_qwen3.5_2b_specdecode.sh - -MODEL_NAME="Qwen/Qwen3.5-2B" -# Use the fast, non-persistent local disk for data (not the ~/default quota). -DATA_DIR="/mnt/local_storage/data/dapo" -TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" -TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" -NUM_NODES=1 -NUM_GPUS_PER_NODE=8 -# 2B is small: TP=1 inference per engine (no TP comm), one engine per GPU. -NUM_INFERENCE_ENGINES=8 -INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 -LOGGER="wandb" # change to "console" to print to stdout - -CLIP_RATIO_LOW=0.2 -CLIP_RATIO_HIGH=0.28 -# use token mean loss reduction -LOSS_REDUCTION="token_mean" -# applies overlong filtering (but not soft overlong punishment) -APPLY_OVERLONG_FILTERING=true -# apply soft overlong punishment with custom trainer impl in main_dapo.py -OVERLONG_BUFFER_LEN=$((1024 * 4)) -OVERLONG_BUFFER_PENALTY_FACTOR=1.0 - -# other DAPO parameters -USE_KL_LOSS=false -TEMPERATURE=1.0 -TOP_P=1.0 -EVAL_TOP_P=0.7 -CLIP_RATIO_C=10.0 -MAX_PROMPT_LENGTH=$((1024 * 2)) -# IDENTICAL to the no-spec run (8192) for an apples-to-apples comparison. The MTP draft loss adds a -# full [seq, vocab] softmax tensor on top of the main lm-head gradient, but this fits at TP=2 with the -# default allocator now that the O(num_microbatches) `mtp_student_logits` accumulation is fixed. -MAX_RESPONSE_LENGTH=$((1024 * 8)) - -# repro run parameters -# Quality-comparison profile: larger batch lowers reward-curve variance (~4x vs batch 32) -# so a spec-decode quality effect is detectable. rollout = 128 * 8 = 1024 seqs. -# MINI_BATCH_SIZE == TRAIN_BATCH_SIZE => 1 on-policy update/rollout. Keep IDENTICAL across no-spec/spec runs. -TRAIN_BATCH_SIZE=128 -MINI_BATCH_SIZE=32 -N_SAMPLES_PER_PROMPT=8 -EVAL_N_SAMPLES_PER_PROMPT=16 -ENFORCE_EAGER=true # cuda graphs can cause some instability -LR=1e-6 - -# megatron config -- Qwen3.5-2B is a dense model, so no expert parallelism. -# TP=2 (matching the no-spec run): the decoupled MTP draft loss materializes a full -# [batch, seq, vocab] float32 softmax tensor on top of the main lm-head gradient at Qwen3.5's 248K -# vocab, but this fits at TP=2 with the default allocator now that the O(num_microbatches) -# `mtp_student_logits` accumulation is fixed. NOTE: raising TP does NOT meaningfully shrink the -# training footprint (it's activation/optimizer-bound, not vocab-bound), and TP=8 also triggered a -# CUDA-IPC weight-sync failure (pidfd_getfd EPERM) when broadcasting the 8-way-sharded policy to the -# TP=1 vLLM engines. So TP=2 is both the apples-to-apples match and the known-good weight-sync config. -# TP=2, PP=1, CP=1 => DP=4. -MEGATRON_TP=2 -MEGATRON_PP=1 -MEGATRON_CP=1 -MEGATRON_EP=1 -MEGATRON_ETP=null - - -# TIS parameters -TIS_IMP_RATIO_CAP=2.0 -TIS_TYPE=token - - -# Multi-Token Prediction (MTP) speculative decoding. -# Qwen3.5-2B ships 1 native MTP head (`mtp_num_hidden_layers: 1`); training always trains the -# checkpoint's heads. NUM_SPECULATIVE_TOKENS is the vLLM *draft depth* only — values > 1 reuse the -# single head autoregressively at draft time (more speedup, but per-position acceptance decays with -# depth since the head never trains on its own outputs). Try 2-3 for extra acceleration. -MTP_ENABLED=true -MTP_NUM_SPECULATIVE_TOKENS=1 -MTP_LOSS_TYPE="soft_ce" # "soft_ce" (distill against policy) | "hard_ce" (ground-truth next tokens) -MTP_LOSS_WEIGHT=0.1 -# NOTE: trainer.policy.megatron_config.mtp_detach_shared_output now defaults to true: the draft loss -# trains ONLY the MTP-head params; the tied embedding/lm_head is detached (output projection AND the -# MTP block's re-embedding), so the draft gradient no longer nudges the policy's own logits. Set it -# to false for the old NeMo-RL behaviour (shared head also trained by the draft loss). -# Top-k draft loss: distill only the teacher's top-k tokens instead of the full 248K vocab, keeping -# draft-loss memory at O(seq*k) vs O(seq*vocab). This is now an optional throughput knob, NOT a -# memory requirement: the full-vocab soft-CE fits comfortably at TP=2 since the real OOM cause (the -# O(num_microbatches) accumulation of `mtp_student_logits`) was fixed by freeing each microbatch's -# student logits right after its draft backward. No PYTORCH_CUDA_ALLOC_CONF tuning is needed. -# null => exact full-vocab soft-CE. Set to e.g. 64 to use the memory-light top-k approximation. -MTP_LOSS_TOPK=null - - -# Qwen3.5 flags -REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 -ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 -DISTRIBUTED_EXECUTOR_BACKEND="mp" -export _SKYRL_USE_NEW_INFERENCE=0 -export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 -# NOTE: do NOT set PYTORCH_CUDA_ALLOC_CONF here (neither expandable_segments:True nor max_split_size_mb). -# - expandable_segments:True makes PyTorch allocate via virtual-memory segments incompatible with the -# legacy CUDA-IPC handle used by colocated weight sync -> it falls back to pidfd_getfd, which this -# cluster's ptrace_scope blocks (pidfd_getfd: Operation not permitted). -# - max_split_size_mb over-reserves PyTorch memory and starves NCCL's external cudaMalloc at grad-sync -# ("Failed to CUDA calloc ... bytes" in reduce_scatter). -# Neither is needed: the draft-loss OOM was an O(num_microbatches) accumulation of `mtp_student_logits`, -# now fixed by freeing each microbatch's student logits right after its draft backward. The full-vocab -# soft-CE fits at TP=2 with the default allocator. - -uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ - data.train_data="['$TRAIN_FILE']" \ - data.val_data="['$TEST_FILE']" \ - trainer.algorithm.advantage_estimator="grpo" \ - trainer.algorithm.policy_loss_type="dual_clip" \ - trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ - trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ - trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ - generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ - generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ - generator.sampling_params.temperature=$TEMPERATURE \ - generator.sampling_params.top_p=$TOP_P \ - generator.eval_sampling_params.top_p=$EVAL_TOP_P \ - generator.eval_sampling_params.temperature=$TEMPERATURE \ - generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ - trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ - trainer.policy.model.path="$MODEL_NAME" \ - trainer.placement.colocate_all=true \ - trainer.strategy=megatron \ - generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ - trainer.placement.policy_num_nodes=$NUM_NODES \ - trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ - generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ - generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ - generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ - trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ - trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ - trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ - trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ - trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ - trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ - trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ - trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ - trainer.epochs=10 \ - trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ - trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ - trainer.eval_batch_size=1024 \ - trainer.eval_before_train=false \ - trainer.eval_interval=5 \ - trainer.update_epochs_per_batch=1 \ - trainer.train_batch_size=$TRAIN_BATCH_SIZE \ - trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ - trainer.micro_forward_batch_size_per_gpu=1 \ - trainer.micro_train_batch_size_per_gpu=1 \ - trainer.ckpt_interval=-1 \ - trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ - generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.policy.optimizer_config.lr=$LR \ - trainer.policy.optimizer_config.num_warmup_steps=5 \ - trainer.policy.optimizer_config.weight_decay=0.1 \ - trainer.policy.optimizer_config.max_grad_norm=1.0 \ - generator.inference_engine.backend=vllm \ - generator.inference_engine.run_engines_locally=true \ - generator.inference_engine.weight_sync_backend=nccl \ - generator.inference_engine.async_engine=true \ - generator.batched=true \ - environment.env_class=aime \ - generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ - generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ - generator.inference_engine.gpu_memory_utilization=0.5 \ - trainer.mtp.enabled=$MTP_ENABLED \ - trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ - trainer.mtp.loss_type=$MTP_LOSS_TYPE \ - trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ - trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ - trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo_sd" \ - trainer.run_name="sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.export_path="/mnt/local_storage/exports/sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.hf_save_interval=300 \ - trainer.resume_mode=latest \ - trainer.max_ckpts_to_keep=3 \ - trainer.ckpt_path="/mnt/local_storage/ckpts/sd_dapo_qwen3_5_2b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - $@ diff --git a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh deleted file mode 100644 index 9b7b0975a7..0000000000 --- a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh +++ /dev/null @@ -1,213 +0,0 @@ -set -x - -# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO with Megatron, -# WITH vLLM MTP speculative decoding for faster rollout, but the MTP head is NOT trained. -# -# This is the "frozen-head" counterpart of run_megatron_dapo_qwen3.5_9b_specdecode.sh. -# Every knob is IDENTICAL except the MTP block at the bottom: the head is BUILT and SYNCED -# like the training variant, but its draft loss weight is 0 so it is never actually trained. -# -# Why not just turn the head off on the training side? Under colocate_all, vLLM is slept -# with level=2 between steps, which DISCARDS its weights -- including the separate spec-decode -# drafter (the MTP head). On wake, SkyRL restores the drafter ONLY from the Megatron -# weight-sync list (inference_servers/spec_decode_utils.py::_reload_spec_decode_drafter). If -# Megatron doesn't build the head, the sync list has no `mtp.*` weights, the drafter stays -# garbage, and draft acceptance is ~0 (many drafted tokens, almost all rejected). So the head -# MUST be built + exported every sync even when we don't want to train it. -# -# What this configuration does: -# - Training side: `trainer.mtp.enabled=true` builds Qwen3.5-9B's native MTP head (the model -# ships `mtp_num_hidden_layers: 1`) and exports it through weight sync each step. But -# `trainer.mtp.loss_weight=0.0` makes the decoupled draft loss contribute ZERO gradient, so -# the head stays at its pretrained values (only ~1e-7/step weight-decay drift). The RL policy -# backbone trains normally. -# - Inference side: vLLM MTP speculative decoding is enabled with `num_speculative_tokens` -# draft tokens; the drafter is re-synced with the (frozen) pretrained head every step. -# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. -# -# CAVEAT: the frozen head is matched to the ORIGINAL base model. As RL updates the policy -# backbone (and syncs those updates to vLLM), the never-updated head drifts out of alignment, -# so `vllm/draft_acceptance_rate` will DECAY over training -- best early, eroding as the policy -# moves. Training the head (the _specdecode.sh variant) closes that gap. -# -# Runs on 1 node of 8xH100s (80GB each). -# -# NOTE: verify the exact HF repo id for the 9B model before running -# (e.g. `hf download Qwen/Qwen3.5-9B` / check https://huggingface.co/Qwen). -# -# Prepare data onto the fast local disk first: -# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh -# Then launch: -# bash examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode_notrain.sh - -MODEL_NAME="Qwen/Qwen3.5-9B" -# Use the fast, non-persistent local disk for data (not the ~/default quota). -DATA_DIR="/mnt/local_storage/data/dapo" -TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" -TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" -NUM_NODES=1 -NUM_GPUS_PER_NODE=8 -# 9B is ~4.5x the 2B: a single full-weight copy is ~18GB in bf16. Use inference -# TP=2 (4 engines) so each engine's weights are halved (~9GB/GPU) and there is -# more headroom for KV cache during generation. TP=2 comm stays on NVLink. -NUM_INFERENCE_ENGINES=8 -INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 -LOGGER="wandb" # change to "console" to print to stdout - -CLIP_RATIO_LOW=0.2 -CLIP_RATIO_HIGH=0.28 -# use token mean loss reduction -LOSS_REDUCTION="token_mean" -# applies overlong filtering (but not soft overlong punishment) -APPLY_OVERLONG_FILTERING=true -# apply soft overlong punishment with custom trainer impl in main_dapo.py -OVERLONG_BUFFER_LEN=$((1024 * 4)) -OVERLONG_BUFFER_PENALTY_FACTOR=1.0 - -# other DAPO parameters -USE_KL_LOSS=false -TEMPERATURE=1.0 -TOP_P=1.0 -EVAL_TOP_P=0.7 -CLIP_RATIO_C=10.0 -MAX_PROMPT_LENGTH=$((1024 * 2)) -MAX_RESPONSE_LENGTH=$((1024 * 8)) - -# repro run parameters -TRAIN_BATCH_SIZE=32 -MINI_BATCH_SIZE=32 -N_SAMPLES_PER_PROMPT=8 -EVAL_N_SAMPLES_PER_PROMPT=16 -ENFORCE_EAGER=true # cuda graphs can cause some instability -LR=1e-6 - -# megatron config -- Qwen3.5-9B is a dense model, so no expert parallelism. -# TP=4 (up from 2 on the 2B): 9B params + Adam states + 8K-token activations -# need more sharding to fit at micro batch 1. TP>1 auto-enables sequence -# parallelism, sharding activations/vocab-logits across the TP group. -# TP=4, PP=1, CP=1 => DP=2. TP stays within the single-node NVLink domain. -MEGATRON_TP=4 -MEGATRON_PP=1 -MEGATRON_CP=1 -MEGATRON_EP=1 -MEGATRON_ETP=null - -# optimizer offload -- OFF by default (with colocate_all the inference engine is -# offloaded during training, so the full 80GB is available for the optimizer). -# Flip to true if you hit OOM in the optimizer step / grad sync. -OPTIMIZER_OFFLOAD=false -OPTIMIZER_OFFLOAD_FRACTION=1.0 - -# TIS parameters -TIS_IMP_RATIO_CAP=2.0 -TIS_TYPE=token - - -# vLLM MTP speculative decoding WITHOUT training the head (frozen pretrained head baseline). -# -# IMPORTANT: under colocate_all, vLLM is slept with level=2 between steps, which DISCARDS its -# weights — including the separate spec-decode drafter (the MTP head). On wake, SkyRL restores the -# drafter ONLY from the Megatron weight-sync list (see inference_servers/spec_decode_utils.py: -# _reload_spec_decode_drafter). So the head MUST be built + exported on the Megatron side every sync, -# or the drafter stays garbage and draft acceptance is ~0 (lots of drafted tokens, almost all -# rejected). vLLM's init-time checkpoint load is NOT enough — sleep(level=2) throws it away on step 1. -# -# Therefore we keep trainer.mtp.enabled=true (builds the native head, exports it each sync, sets -# vLLM speculative_config) but set loss_weight=0.0 so the draft loss contributes ZERO gradient: the -# head stays at its pretrained values (only ~1e-7/step weight-decay drift, negligible) and is -# re-synced to vLLM every step. Acceptance then reflects the genuine "stock pretrained head, drifting -# RL backbone" baseline. NOTE: the (top-k soft-CE) draft loss is still computed each step and thrown -# away — small wasted compute. Ask for the proper freeze path (skip the loss + freeze the params) if -# that matters. -MTP_NUM_SPECULATIVE_TOKENS=3 -MTP_LOSS_TYPE="soft_ce" # kept cheap via top-k; weight is 0 so the value is discarded -MTP_LOSS_WEIGHT=0.0 # 0 => head is built + synced but never trained (frozen pretrained head) -MTP_LOSS_TOPK=256 - - -# Qwen3.5 flags -REMOVE_MICROBATCH_PADDING=false # sample packing is not yet supported for GDN layers in megatron - see: https://github.com/NVIDIA/Megatron-LM/pull/2644 -ENGINE_INIT_KWARGS='{"gdn_prefill_backend": "triton"}' # see https://github.com/vllm-project/vllm/issues/36921#issuecomment-4109702738 -DISTRIBUTED_EXECUTOR_BACKEND="mp" -export _SKYRL_USE_NEW_INFERENCE=0 -export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 - -uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ - data.train_data="['$TRAIN_FILE']" \ - data.val_data="['$TEST_FILE']" \ - trainer.algorithm.advantage_estimator="grpo" \ - trainer.algorithm.policy_loss_type="dual_clip" \ - trainer.algorithm.overlong_buffer_len=$OVERLONG_BUFFER_LEN \ - trainer.algorithm.overlong_buffer_penalty_factor=$OVERLONG_BUFFER_PENALTY_FACTOR \ - trainer.algorithm.loss_reduction=$LOSS_REDUCTION \ - generator.inference_engine.enforce_eager=$ENFORCE_EAGER \ - generator.apply_overlong_filtering=$APPLY_OVERLONG_FILTERING \ - generator.sampling_params.temperature=$TEMPERATURE \ - generator.sampling_params.top_p=$TOP_P \ - generator.eval_sampling_params.top_p=$EVAL_TOP_P \ - generator.eval_sampling_params.temperature=$TEMPERATURE \ - generator.eval_sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.algorithm.use_kl_loss=$USE_KL_LOSS \ - trainer.algorithm.clip_ratio_c=$CLIP_RATIO_C \ - trainer.policy.model.path="$MODEL_NAME" \ - trainer.placement.colocate_all=true \ - trainer.strategy=megatron \ - generator.inference_engine.distributed_executor_backend="$DISTRIBUTED_EXECUTOR_BACKEND" \ - trainer.placement.policy_num_nodes=$NUM_NODES \ - trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ - generator.inference_engine.engine_init_kwargs="$ENGINE_INIT_KWARGS" \ - generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ - generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ - trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ - trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ - trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ - trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ - trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ - trainer.policy.megatron_config.optimizer_config_kwargs.overlap_cpu_optimizer_d2h_h2d=$OPTIMIZER_OFFLOAD \ - trainer.policy.megatron_config.optimizer_config_kwargs.use_precision_aware_optimizer=$OPTIMIZER_OFFLOAD \ - trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_cpu_offload=$OPTIMIZER_OFFLOAD \ - trainer.policy.megatron_config.optimizer_config_kwargs.optimizer_offload_fraction=$OPTIMIZER_OFFLOAD_FRACTION \ - trainer.algorithm.off_policy_correction.tis_ratio_type=$TIS_TYPE \ - trainer.algorithm.off_policy_correction.token_tis_ratio_clip_high=$TIS_IMP_RATIO_CAP \ - trainer.remove_microbatch_padding=$REMOVE_MICROBATCH_PADDING \ - trainer.epochs=10 \ - trainer.algorithm.eps_clip_low=$CLIP_RATIO_LOW \ - trainer.algorithm.eps_clip_high=$CLIP_RATIO_HIGH \ - trainer.eval_batch_size=1024 \ - trainer.eval_before_train=false \ - trainer.eval_interval=5 \ - trainer.update_epochs_per_batch=1 \ - trainer.train_batch_size=$TRAIN_BATCH_SIZE \ - trainer.policy_mini_batch_size=$MINI_BATCH_SIZE \ - trainer.micro_forward_batch_size_per_gpu=1 \ - trainer.micro_train_batch_size_per_gpu=1 \ - trainer.ckpt_interval=50 \ - trainer.max_prompt_length=$MAX_PROMPT_LENGTH \ - generator.sampling_params.max_generate_length=$MAX_RESPONSE_LENGTH \ - trainer.policy.optimizer_config.lr=$LR \ - trainer.policy.optimizer_config.num_warmup_steps=5 \ - trainer.policy.optimizer_config.weight_decay=0.1 \ - trainer.policy.optimizer_config.max_grad_norm=1.0 \ - generator.inference_engine.backend=vllm \ - generator.inference_engine.run_engines_locally=true \ - generator.inference_engine.weight_sync_backend=nccl \ - generator.inference_engine.async_engine=true \ - generator.batched=true \ - environment.env_class=aime \ - generator.n_samples_per_prompt=$N_SAMPLES_PER_PROMPT \ - generator.eval_n_samples_per_prompt=$EVAL_N_SAMPLES_PER_PROMPT \ - generator.inference_engine.gpu_memory_utilization=0.5 \ - trainer.mtp.enabled=true \ - trainer.mtp.num_speculative_tokens=$MTP_NUM_SPECULATIVE_TOKENS \ - trainer.mtp.loss_type=$MTP_LOSS_TYPE \ - trainer.mtp.loss_weight=$MTP_LOSS_WEIGHT \ - trainer.policy.megatron_config.mtp_loss_topk=$MTP_LOSS_TOPK \ - trainer.logger="$LOGGER" \ - trainer.project_name="qwen3_5_dapo_2" \ - trainer.run_name="sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.export_path="/mnt/local_storage/exports/sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - trainer.hf_save_interval=300 \ - trainer.resume_mode=latest \ - trainer.max_ckpts_to_keep=3 \ - trainer.ckpt_path="/mnt/local_storage/ckpts/sd_notrain_dapo_qwen3_5_9b_megatron_tp${MEGATRON_TP}_pp${MEGATRON_PP}_cp${MEGATRON_CP}" \ - $@ diff --git a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py index 77be1263cd..397bae8803 100644 --- a/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py +++ b/skyrl/backends/skyrl_train/inference_engines/vllm/spec_decode_metrics.py @@ -1,19 +1,14 @@ -# Cumulative speculative-decoding (MTP draft) acceptance accounting for vLLM v1. -# -# vLLM's async engine (the one SkyRL uses for generation + weight sync) does not expose -# ``get_metrics()``, so we attach a tiny custom stat logger that runs in the AsyncLLM frontend -# process (= the SkyRL engine actor) and accumulates the per-iteration spec-decode counters -# (``num_draft_tokens`` / ``num_accepted_tokens`` / ``num_accepted_tokens_per_pos``). The engine -# reads these cumulative counters and the generator turns them into per-step acceptance rates: -# one overall rate (accepted / drafted) plus one rate per draft position (how often the k-th -# speculated token was accepted), so multi-token drafting (num_speculative_tokens > 1) shows the -# per-depth acceptance decay. The number of per-position metrics tracks whatever depth vLLM -# reports, so it grows/shrinks with the configured draft depth automatically. +# Cumulative speculative-decoding (MTP draft) acceptance accounting for vLLM v1. vLLM's async engine +# has no ``get_metrics()``, so we attach a tiny stat logger in the AsyncLLM frontend that accumulates +# the per-iteration draft/accept counters. The generator turns the cumulative counts into per-step +# acceptance rates -- one overall and one per draft position (grows with num_speculative_tokens). from __future__ import annotations from typing import Any, Optional +from loguru import logger + def _add_per_pos(into: list, add) -> None: """Elementwise-add a per-position count list into ``into``, growing ``into`` as needed.""" @@ -35,6 +30,9 @@ def make_spec_decode_stat_logger_class(): class SpecDecodeStatLogger(StatLoggerBase): """Accumulate cumulative spec-decode draft/accept counts across scheduler iterations.""" + # Process-global one-shot guard: warn once if vLLM stops exposing `spec_decoding_stats`. + _warned_missing_spec_field = False + def __init__(self, vllm_config, engine_index: int = 0): self.engine_index = engine_index self.num_drafts = 0 @@ -46,7 +44,20 @@ def __init__(self, vllm_config, engine_index: int = 0): self.num_accepted_tokens_per_pos: list = [] def record(self, scheduler_stats=None, iteration_stats=None, mm_cache_stats=None, engine_idx: int = 0): - stats = getattr(scheduler_stats, "spec_decoding_stats", None) if scheduler_stats is not None else None + if scheduler_stats is None: + return + # Absent field (vs. present-but-None = no drafting this iter) means vLLM renamed it; warn + # once so the metrics don't silently go empty. + if not hasattr(scheduler_stats, "spec_decoding_stats"): + if not SpecDecodeStatLogger._warned_missing_spec_field: + SpecDecodeStatLogger._warned_missing_spec_field = True + logger.warning( + "vLLM SchedulerStats has no 'spec_decoding_stats' field; spec-decode (MTP " + "draft) acceptance metrics will be empty. vLLM's metrics API likely changed " + "-- update inference_engines/vllm/spec_decode_metrics.py." + ) + return + stats = scheduler_stats.spec_decoding_stats if stats is not None: self.num_drafts += stats.num_drafts self.num_draft_tokens += stats.num_draft_tokens diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index 34231df827..12cc00931b 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -143,17 +143,12 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: else: args.enable_lora = False - # Speculative decoding (e.g. Multi-Token Prediction). Pass the dict straight through to vLLM's - # AsyncEngineArgs.speculative_config. For MTP (`{"method": "mtp", ...}`) vLLM loads the MTP heads - # from the same policy checkpoint, so SkyRL's weight sync keeps the draft aligned with the policy - # (the trained MTP head weights are exported by the Megatron bridge and loaded into the draft). + # Speculative decoding (e.g. MTP): passed straight through to vLLM's speculative_config. if ie_cfg.speculative_config is not None: spec_cfg = get_config_as_dict(ie_cfg.speculative_config) args.speculative_config = spec_cfg logger.info(f"vLLM speculative decoding enabled: speculative_config={spec_cfg}") - # Add any extra engine_init_kwargs (applied last so they can override anything above, including - # speculative_config, if a user needs full control). engine_kwargs = get_config_as_dict(ie_cfg.engine_init_kwargs) for key, value in engine_kwargs.items(): setattr(args, key, value) diff --git a/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py b/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py index 70d4c29c0b..1ed5e5e9a0 100644 --- a/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py +++ b/skyrl/backends/skyrl_train/inference_servers/vllm_worker.py @@ -115,8 +115,7 @@ def load_weights(self, request: bytes) -> None: else: self.model_runner.reload_weights(weights_iterator=iter(weight_list)) # Re-sync the spec-decode (MTP/Eagle) drafter from the same weights after the main load; - # the main-model load never touches the separate drafter module. No-op when spec decoding - # is disabled. + # the main-model load never touches the separate drafter module. _reload_spec_decode_drafter(self.model_runner, weight_list) if weight_update_bracketed: diff --git a/skyrl/backends/skyrl_train/mtp/hidden_capture.py b/skyrl/backends/skyrl_train/mtp/hidden_capture.py index 633f019569..ce2c5e5ec2 100644 --- a/skyrl/backends/skyrl_train/mtp/hidden_capture.py +++ b/skyrl/backends/skyrl_train/mtp/hidden_capture.py @@ -1,16 +1,8 @@ -# Decoupled capture of the trunk hidden states that feed the MTP/draft head. Inspired by NeMo-RL's -# draft-model training (https://github.com/NVIDIA-NeMo/RL), adapted for SkyRL's Megatron backend. -# -# Works for any Megatron model with native MTP heads exposing ``self.mtp`` (the shared -# ``MultiTokenPredictionBlock``): ``GPTModel`` (DeepSeek-V3 / GLM / Qwen3-Next) and ``MambaModel`` -# (Qwen3.5 / NemotronH). The in-forward ``self.mtp(...)`` passes the trunk hidden states through to -# ``process_mtp_loss`` as the main logits, so we must not tamper with it. Instead: -# 1. a forward pre-hook records the kwargs Megatron built for ``self.mtp`` (without modifying them); -# 2. after the forward, we re-invoke ``self.mtp`` with those kwargs but the trunk hidden states -# detached -- the decoupled draft pass, whose gradient reaches the MTP head but never the policy -# backbone. Reusing the kwargs avoids reconstructing rotary embeddings / attention masks. -# The replay costs one extra MTP-block forward (typically a single layer) -- a deliberate trade for -# robustness over reaching into GPTModel's rotary/mask construction. +# Decoupled capture of the trunk hidden states feeding the MTP/draft head (inspired by NeMo-RL, +# adapted for SkyRL's Megatron backend). Works for any Megatron model exposing ``self.mtp``. +# A forward pre-hook records the kwargs Megatron builds for ``self.mtp``; after the forward we +# re-invoke it with the trunk hidden states detached, so the draft gradient reaches the head but not +# the policy backbone. Reusing the captured kwargs avoids rebuilding rotary embeddings / masks. from __future__ import annotations @@ -50,15 +42,15 @@ def _mtp_layer_offset(mtp_block) -> int: (the trunk hidden states plus any MTP outputs forwarded from earlier pipeline/VP stages), appends one new chunk per MTP depth, and concatenates everything along dim 0. ``offset`` is 0 in the common single-stage case; we resolve it via megatron-core so the replay split below is also - correct under PP/VPP. Falls back to 0 if the helper is unavailable (older megatron-core).""" + correct under PP/VPP. Falls back to 0 only if the helper is missing (older megatron-core); other + errors propagate rather than silently returning 0 and mis-splitting the hidden states.""" try: from megatron.core.transformer.multi_token_prediction import ( get_mtp_layer_offset, ) - - return int(get_mtp_layer_offset(mtp_block.config, getattr(mtp_block, "vp_stage", None))) - except Exception: + except ImportError: return 0 + return int(get_mtp_layer_offset(mtp_block.config, getattr(mtp_block, "vp_stage", None))) class MTPHiddenCapture: diff --git a/skyrl/backends/skyrl_train/mtp/separate_optim.py b/skyrl/backends/skyrl_train/mtp/mtp_optim.py similarity index 75% rename from skyrl/backends/skyrl_train/mtp/separate_optim.py rename to skyrl/backends/skyrl_train/mtp/mtp_optim.py index 971b8db65a..5bd30c0413 100644 --- a/skyrl/backends/skyrl_train/mtp/separate_optim.py +++ b/skyrl/backends/skyrl_train/mtp/mtp_optim.py @@ -1,42 +1,8 @@ -"""C-full: isolate the MTP / draft head into its OWN grad buffer + ``DistributedOptimizer``. - -Why this exists ---------------- -The decoupled draft loss is autograd-clean (it only reaches ``.mtp.*`` params; the trunk hidden and -the shared output/embedding weights are detached). But when the MTP head shares the policy's single -Megatron DDP grad buffer, the head's mere *presence* changes the layout of that buffer and therefore -the floating-point result of the POLICY gradient's distributed reduction (the DP reduce-scatter and the -TP all-reduce of sequence-parallel / layernorm grads, both of which coalesce every ``requires_grad`` -param into one flat tensor). That perturbation is deterministic and systematic, and the RL feedback -loop amplifies it into the policy-entropy collapse observed on MiMo-7B with MTP enabled. - -Neither separating the grad *clip* (the head's grad does not dilute the policy via the clip — verified) -nor pinning the NCCL algorithm fixes it, because the channel is the shared *buffer/reduction* itself. -The robust fix is to make the policy's grad buffer + reduction byte-identical to a model built with NO -MTP head, while still co-training the head at full strength. - -Mechanism ---------- -Megatron exposes no "exclude these params from the buffer" API; the only lever is ``requires_grad`` at -DDP-construction time (DDP buckets only ``requires_grad`` params, and changing it *after* the wrap -deadlocks distributed-optimizer init). So: - -1. :func:`freeze_mtp_params_pre_wrap` (a provider pre-wrap hook) sets the ``.mtp`` params - ``requires_grad=False`` *before* the policy is DDP-wrapped, so they are excluded from the policy - grad buffer -> the policy buffer is byte-identical to the no-MTP build. -2. After the policy optimizer is built, :class:`SeparateMTPOptimizer` flips the head back to - ``requires_grad=True`` and wraps the SAME ``host.mtp`` submodule (left in place inside ``GPTModel``, - so capture / replay / weight-export are unchanged) in its own Megatron DDP + ``DistributedOptimizer``. -3. The policy's own ``finalize_model_grads`` would still coalesce the head's layernorm grads into the - policy TP all-reduce (it iterates every ``requires_grad`` param of the GPTModel). :func:`make_policy_finalize_excluding_mtp` - wraps it to transiently hide the head during the policy finalize (backward is already complete, so - this is only a read-time filter), keeping the policy reduction byte-identical to the no-MTP build. -4. The combined ``policy + weight * draft`` backward auto-routes each param's grad to its own buffer - (the draft loss is autograd-decoupled). We finalize + step the head's optimizer alongside the policy's. - -Acceptance test: with C-full on, the fixed-batch no-clip policy-param trajectory must match the no-MTP -run within the run-to-run nondeterminism floor (see -``tests/.../megatron/test_mtp_grad_coupling.py::test_mtp_fixed_batch_param_drift``). +"""Isolate the MTP / draft head into its OWN grad buffer + ``DistributedOptimizer``. + +Even though the draft loss is autograd-clean (it only touches ``.mtp.*`` params), letting the head +share the policy's single Megatron DDP grad buffer changes that buffer's layout and perturbs the +floating-point result of the policy gradient's distributed reduction. """ from __future__ import annotations @@ -78,7 +44,7 @@ def freeze_mtp_params_pre_wrap(model_or_models: Union[torch.nn.Module, List[torc """Provider pre-wrap hook: set the MTP/draft-head params ``requires_grad=False`` so Megatron's DDP excludes them from the POLICY grad buffer (the only Megatron lever for buffer exclusion is ``requires_grad`` at construction time). They are re-enabled and given their own buffer + optimizer - by :class:`SeparateMTPOptimizer` after the policy is wrapped. Result: the policy grad buffer and its + by :class:`MTPOptimizer` after the policy is wrapped. Result: the policy grad buffer and its distributed reduction are byte-identical to a model built with no MTP head. Modifies in place AND returns the model unchanged. (The bridge's ``pre_wrap_hook`` property @@ -86,10 +52,20 @@ def freeze_mtp_params_pre_wrap(model_or_models: Union[torch.nn.Module, List[torc break the chain for any subsequent hook — we must return the model.) """ models = model_or_models if isinstance(model_or_models, list) else [model_or_models] + num_frozen = 0 for model in models: for name, param in model.named_parameters(): if is_mtp_param_name(name): param.requires_grad = False + num_frozen += 1 + # Only registered when C-full is on, so zero matches means the `.mtp.` naming drifted -- the head + # would silently train in the policy grad buffer and defeat the isolation. Fail loud instead. + if num_frozen == 0: + raise RuntimeError( + "freeze_mtp_params_pre_wrap matched no parameters via is_mtp_param_name (expected '.mtp.' " + "or 'mtp.' in the name). The MTP submodule naming likely changed -- C-full grad isolation " + "would silently break. Update is_mtp_param_name in mtp/mtp_optim.py." + ) return model_or_models @@ -101,7 +77,7 @@ def make_policy_finalize_excluding_mtp(mtp_params: List[torch.nn.Parameter]): head's grads were included, that coalesced reduction would differ from the no-MTP build -> the exact perturbation we are eliminating. Backward is already complete when finalize runs, so transiently clearing the head's ``requires_grad`` is a pure read-time filter (no effect on the head's grads, - which already live in the separate MTP buffer and are reduced by :meth:`SeparateMTPOptimizer.step`). + which already live in the separate MTP buffer and are reduced by :meth:`MTPOptimizer.step`). """ def finalize(model, *args, **kwargs): @@ -121,7 +97,7 @@ def _build_mtp_ddp_config(ddp_config) -> DistributedDataParallelConfig: """Build the head's ``DistributedDataParallelConfig`` from the policy's, but with overlap disabled. ``overlap_grad_reduce=False`` so the head's grads simply accumulate into its buffer across - micro-batches and are reduced once in :meth:`SeparateMTPOptimizer.finalize_grads` — there is no + micro-batches and are reduced once in :meth:`MTPOptimizer.finalize_grads` — there is no per-microbatch sync to coordinate with the policy's pipeline schedule (the head is not in the chunk list passed to ``forward_backward_func``). ``overlap_param_gather=False`` for the same reason. """ @@ -135,7 +111,7 @@ def _build_mtp_ddp_config(ddp_config) -> DistributedDataParallelConfig: return cfg -class SeparateMTPOptimizer: +class MTPOptimizer: """Owns the MTP/draft head's isolated grad buffer + ``DistributedOptimizer`` (C-full). The head stays inside the policy ``GPTModel`` (so capture / replay / weight-export are unchanged) but @@ -153,7 +129,7 @@ def __init__(self, policy_module, ddp_config, optim_config, scheduler_config, nu self.mtp_module = _resolve_mtp_module(policy_module) if self.mtp_module is None: raise RuntimeError( - "SeparateMTPOptimizer: no `.mtp` head found on the policy model. Enable MTP " + "MTPOptimizer: no `.mtp` head found on the policy model. Enable MTP " "(trainer.mtp.enabled / mtp_num_layers) or disable mtp_separate_optimizer." ) diff --git a/skyrl/backends/skyrl_train/mtp/native_loss_patch.py b/skyrl/backends/skyrl_train/mtp/native_loss_patch.py index 57024f9c6b..96ecd4809d 100644 --- a/skyrl/backends/skyrl_train/mtp/native_loss_patch.py +++ b/skyrl/backends/skyrl_train/mtp/native_loss_patch.py @@ -1,26 +1,13 @@ """Disable Megatron's native in-forward MTP loss (``process_mtp_loss``). -SkyRL trains the MTP/draft head with its OWN decoupled soft-CE distillation loss (see -``mtp/soft_ce.py`` and ``MegatronModelWrapper``), NOT Megatron's pretraining-style MTP auxiliary -loss. But ``GPTModel.forward`` / ``HybridModel.forward`` call ``process_mtp_loss`` unconditionally -whenever the model is built with MTP heads (``mtp_num_layers`` set) and run in training/eval. That -native loss is a hard cross-entropy whose gradient — unless ``mtp_detach_heads`` is set — flows -straight into the shared trunk and output embedding, corrupting the RL policy (inflated grad-norm -and entropy collapse), with a magnitude set by Megatron's own scaling (so it is independent of -SkyRL's ``mtp_loss_weight``). +SkyRL trains the MTP head with its own decoupled soft-CE loss, but ``GPTModel``/``HybridModel`` call +``process_mtp_loss`` unconditionally when MTP heads exist; its hard-CE gradient flows into the shared +trunk and corrupts the RL policy (inflated grad-norm, entropy collapse). SkyRL used to short-circuit +it by passing no labels, but a megatron-core update derives labels from ``input_ids`` -- so we now +replace ``process_mtp_loss`` at its call sites with a no-op instead. -Megatron only short-circuits the native loss when BOTH ``labels`` and ``input_ids`` are None. A -megatron-core update added "derive labels from ``input_ids`` (e.g. RL training)", so passing no -labels no longer disables it. Rather than depend on that label / forward-gating behaviour (which is -exactly what silently broke), we replace ``process_mtp_loss`` at its call sites with a no-op that -returns the main-model hidden chunk and computes no loss. SkyRL's decoupled head training is -unaffected — the head still trains via SkyRL's explicit loss. - -Robustness: this patch is loud, not silent. It raises if a target module no longer exposes -``process_mtp_loss`` (a Megatron rename). Its one blind spot is Megatron *inlining* the loss into -the forward instead of calling the function — which the grad-isolation acceptance test (the policy -gradient must match a no-MTP build, see ``tests/.../megatron/test_mtp_grad_coupling.py``) is there -to catch. +Loud, not silent: raises if a target module no longer exposes ``process_mtp_loss`` (a Megatron +rename). Blind spot: Megatron inlining the loss into the forward (no test currently covers this). """ from __future__ import annotations diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 2fecf78272..691cb2a2d1 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -864,7 +864,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): # MTP C-full: isolate the draft head into its own grad buffer + optimizer so the policy's # distributed grad reduction is byte-identical to a no-MTP model (the head's presence in the # shared grad buffer otherwise perturbs the policy reduction and collapses entropy — see - # mtp/separate_optim.py). Step 1, here: freeze the head BEFORE the policy DDP wrap so Megatron's + # mtp/mtp_optim.py). Step 1, here: freeze the head BEFORE the policy DDP wrap so Megatron's # DDP (which only buckets requires_grad params) excludes it from the policy grad buffer. The head # is re-enabled and given its own buffer + optimizer after the policy optimizer is built. self._mtp_separate = None @@ -872,7 +872,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): self.cfg.policy.megatron_config.mtp_separate_optimizer and getattr(self.provider, "mtp_num_layers", None) ) if self._mtp_cfull_enabled: - from skyrl.backends.skyrl_train.mtp.separate_optim import ( + from skyrl.backends.skyrl_train.mtp.mtp_optim import ( freeze_mtp_params_pre_wrap, ) @@ -925,17 +925,17 @@ def init_model(self, model_path, num_training_steps: int = 1e9): # MTP C-full step 2: re-enable the draft head and give it its OWN grad buffer + # DistributedOptimizer (the policy optimizer above already excludes it — it was frozen # before the policy wrap). The head co-trains at full strength while the policy reduction - # stays byte-identical to a no-MTP model. See mtp/separate_optim.py. + # stays byte-identical to a no-MTP model. See mtp/mtp_optim.py. if self._mtp_cfull_enabled: - from skyrl.backends.skyrl_train.mtp.separate_optim import ( - SeparateMTPOptimizer, + from skyrl.backends.skyrl_train.mtp.mtp_optim import ( + MTPOptimizer, ) # Fresh optim config for the head (don't share the policy's object). mtp_optim_config = init_megatron_optim_config( self.cfg.policy.optimizer_config, self.cfg.policy.megatron_config.optimizer_config_kwargs ) - self._mtp_separate = SeparateMTPOptimizer( + self._mtp_separate = MTPOptimizer( policy_module=self.actor_module[0], ddp_config=self.cfg.policy.megatron_config.ddp_config, optim_config=mtp_optim_config, @@ -960,7 +960,7 @@ def init_model(self, model_path, num_training_steps: int = 1e9): if self._mtp_separate is not None: from megatron.core.utils import get_model_config - from skyrl.backends.skyrl_train.mtp.separate_optim import ( + from skyrl.backends.skyrl_train.mtp.mtp_optim import ( make_policy_finalize_excluding_mtp, ) diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 06603a928e..b871a9875c 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -249,7 +249,7 @@ class MegatronConfig(BaseConfig): grads). That perturbation is deterministic and the RL feedback loop amplifies it into policy-entropy collapse. With this on, the policy's grad buffer and reduction are byte-identical to a model built with no MTP head, while the head still co-trains at full strength. Only used when MTP heads are - active. See skyrl/backends/skyrl_train/mtp/separate_optim.py.""" + active. See skyrl/backends/skyrl_train/mtp/mtp_optim.py.""" mtp_loss_chunk_size: Optional[int] = 1024 """Sequence-chunk size for the draft loss, with gradient checkpointing, to bound peak memory at large vocab (e.g. Qwen3.5's 248K, where the full-sequence softmax OOMs). Numerically identical to diff --git a/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py b/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py index d3f0bdc760..a65b5c3a38 100644 --- a/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py +++ b/tests/backends/skyrl_train/mtp/test_build_vllm_cli_args_mtp.py @@ -12,6 +12,7 @@ @pytest.mark.vllm def test_build_vllm_cli_args_speculative_config_mtp(monkeypatch): """speculative_config is passed through to vLLM for MTP draft decoding.""" + pytest.importorskip("vllm") import vllm.platforms from vllm.platforms.interface import UnspecifiedPlatform From f69e38a81e7c9dab612661fa0dc168b57db04608 Mon Sep 17 00:00:00 2001 From: zanderjiang Date: Thu, 25 Jun 2026 21:30:44 +0000 Subject: [PATCH 28/28] simpliy commentation --- .ray_cache_bust | 1 - ...run_megatron_dapo_mimo_7b_rl_specdecode.sh | 44 +-------- ...run_megatron_dapo_qwen3.5_9b_specdecode.sh | 33 +------ .../run_search_megatron_mimo_7b_specdecode.sh | 23 +---- skyrl/backends/skyrl_train/mtp/mtp_optim.py | 10 +- .../workers/megatron/megatron_worker.py | 93 +++++++------------ skyrl/train/config/config.py | 2 +- skyrl/train/utils/utils.py | 24 ++--- 8 files changed, 55 insertions(+), 175 deletions(-) delete mode 100644 .ray_cache_bust diff --git a/.ray_cache_bust b/.ray_cache_bust deleted file mode 100644 index 4a78c7e881..0000000000 --- a/.ray_cache_bust +++ /dev/null @@ -1 +0,0 @@ -1780815368 diff --git a/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh index 481d5e9deb..58aeb1325e 100644 --- a/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh +++ b/examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh @@ -1,43 +1,9 @@ set -x -# Colocated DAPO training+generation for XiaomiMiMo/MiMo-7B-RL (dense) on DAPO with -# Megatron, WITH native Multi-Token Prediction (MTP) speculative decoding for faster rollout. -# -# This is the MiMo-7B-RL counterpart of run_megatron_dapo_qwen3.5_9b_specdecode.sh. The -# DAPO knobs (batch sizes, LR, sampling, clipping) are kept identical to the Qwen recipe so the -# reward-curve / throughput comparison stays apples-to-apples; only the model and the few -# architecture-specific flags below differ. -# -# Why MiMo-7B-RL is a good MTP target: -# - MiMo was pretrained WITH MTP. The released checkpoint ships one native MTP module -# (`num_nextn_predict_layers: 1`, the DeepSeek-style key), so both the training-side head and -# the inference-side drafter come from the same pretrained weights — no cold-start head. -# - The backbone is a plain Qwen2-style dense transformer (36 layers, full attention, -# `attention_bias: true`), so there are NO GDN / linear-attention layers. Two consequences vs -# the Qwen3.5 recipe: (a) megatron sample packing IS supported, so we set -# remove_microbatch_padding=true for better throughput; (b) we drop the `gdn_prefill_backend` -# vLLM engine kwarg — it does nothing here. -# - Embeddings are UNTIED (`tie_word_embeddings: false`). The decoupled draft loss therefore -# isolates the output projection by detaching `output_layer.weight` (the untied branch of -# mtp_detach_shared_output, on by default), so the draft gradient trains ONLY the MTP-head params. -# -# What MTP on does (single high-level `trainer.mtp` knob, see skyrl/train/config/config.py): -# - Training side: megatron-bridge's MiMo bridge builds MiMo-7B-RL's native MTP head (from -# `num_nextn_predict_layers`) and trains it with a decoupled draft loss (top-k soft-CE -# distillation against the policy's own detached next-token distribution). The draft gradient -# never pulls on the policy backbone. -# - Inference side: enables vLLM MTP speculative decoding. The generic `{"method": "mtp"}` config -# auto-detects MiMo's `mimo_mtp` drafter from the model config. With a single trained head, -# depth>1 reuses that head autoregressively at draft time. SkyRL's weight sync keeps the draft -# head in sync with the trained policy each step. -# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. -# -# Runs on 1 node of 8xH100s (80GB each). -# -# Prepare data onto the fast local disk first: -# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh -# Then launch: -# bash examples/train/megatron/run_megatron_dapo_mimo_7b_rl_specdecode.sh +# Colocated DAPO training+generation for XiaomiMiMo/MiMo-7B-RL (dense) on DAPO data with Megatron, +# with native Multi-Token Prediction (MTP) speculative decoding for faster rollout. Runs on 1x8 H100. +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/spec_decode/run_megatron_dapo_mimo_7b_rl_specdecode.sh MODEL_NAME="XiaomiMiMo/MiMo-7B-RL" # Use the fast, non-persistent local disk for data (not the ~/default quota). @@ -122,7 +88,7 @@ MTP_LOSS_WEIGHT=0.2 # Top-k draft loss: distill only the teacher's top-k tokens instead of the full vocab, keeping # draft-loss memory at O(seq*k) vs O(seq*vocab). Only applies to mtp_loss_type="soft_ce". Here k=256. MTP_LOSS_TOPK=256 -# Train the MTP head in the policy's native DDP buffer + optimizer (no separate C-full buffer); the +# Train the MTP head in the policy's native DDP buffer + optimizer (no separate MTP buffer); the # native MTP loss stays patched off and the draft loss stays autograd-decoupled, so no MTP grad # reaches the policy -- only the ~ulp-level (DP>=3) grad-buffer reduction reordering is given up. MTP_SEPARATE_OPTIMIZER=false diff --git a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh index bc134045ce..8755555dce 100644 --- a/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh +++ b/examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh @@ -1,34 +1,9 @@ set -x -# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO with Megatron, -# WITH Multi-Token Prediction (MTP) speculative decoding for faster rollout. -# -# This is the spec-decode counterpart of run_megatron_dapo_qwen3.5_9b.sh. Every -# knob below is IDENTICAL to the no-spec script (same batch sizes, LR, parallelism, -# sampling) so a reward-curve / throughput comparison is apples-to-apples. The ONLY -# difference is the `trainer.mtp.*` block at the bottom. -# -# What MTP on does (single high-level `trainer.mtp` knob, see skyrl/train/config/config.py): -# - Training side: builds + trains Qwen3.5-9B's native MTP head (the model ships -# `mtp_num_hidden_layers: 1`) with a decoupled draft loss (top-k soft-CE distillation against -# the policy's own detached next-token distribution). The draft gradient never pulls on -# the policy backbone. -# - Inference side: enables vLLM MTP speculative decoding -# (`speculative_config={"method": "mtp", "num_speculative_tokens": 3}`). With a single trained -# head, depth>1 reuses that head autoregressively at draft time. vLLM loads the MTP head from the -# same policy checkpoint, and SkyRL's weight sync keeps the draft head in sync with the trained -# policy each step. -# - The per-step draft acceptance rate is logged as `vllm/draft_acceptance_rate`. -# -# Runs on 1 node of 8xH100s (80GB each). -# -# NOTE: verify the exact HF repo id for the 9B model before running -# (e.g. `hf download Qwen/Qwen3.5-9B` / check https://huggingface.co/Qwen). -# -# Prepare data onto the fast local disk first: -# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh -# Then launch: -# bash examples/train/megatron/run_megatron_dapo_qwen3.5_9b_specdecode.sh +# Colocated DAPO training+generation for Qwen3.5-9B (dense) on DAPO data with Megatron, +# with Multi-Token Prediction (MTP) speculative decoding for faster rollout. Runs on 1x8 H100. +# DATA_DIR=/mnt/local_storage/data/dapo bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/spec_decode/run_megatron_dapo_qwen3.5_9b_specdecode.sh MODEL_NAME="Qwen/Qwen3.5-9B" # Use the fast, non-persistent local disk for data (not the ~/default quota). diff --git a/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh b/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh index 3719ab7273..e0f079b479 100644 --- a/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh +++ b/examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh @@ -1,34 +1,15 @@ set -x # Colocated GRPO multi-turn SearchR1 training+generation for XiaomiMiMo/MiMo-7B-RL (dense) with -# Megatron, WITH native Multi-Token Prediction (MTP) speculative decoding (k=3) for faster rollout. -# -# This is the MiMo-7B-RL + Megatron + spec-decode counterpart of examples/train/search/run_search.sh -# (which trains Qwen2.5-3B-Instruct with FSDP and NO spec decode). The SearchR1 recipe knobs -# (dataset, search env, multi-turn sampling, search server) are kept identical to run_search.sh / -# run_search_megatron.sh; only the model, the FSDP->Megatron strategy, and the MTP block differ. -# -# MTP on MiMo-7B-RL (see run_megatron_dapo_mimo_7b_rl_specdecode.sh for the full rationale): -# - MiMo ships 1 native MTP head (`num_nextn_predict_layers: 1`), so both the training-side head -# and the inference-side drafter come from the same pretrained weights (no cold-start head). -# - Plain Qwen2-style dense backbone (no GDN), so megatron sample packing IS supported -# (remove_microbatch_padding=true) and no special vLLM prefill backend is needed. -# - Untied embeddings => the decoupled draft loss (mtp_detach_shared_output default true) trains -# ONLY the MTP-head params; the draft gradient never nudges the policy backbone. -# - vLLM MTP speculative decoding draft depth = k = num_speculative_tokens. Here k=3: the single -# trained head is reused autoregressively at draft time. Acceptance is logged as -# `vllm/draft_acceptance_rate`. -# -# Runs on 1 node of 8xH100s (80GB each). +# Megatron, with native Multi-Token Prediction (MTP) speculative decoding (k=3). Runs on 1x8 H100. # # Setup (see examples/train/search/README.md): # 1. Dataset parquet: # uv run --isolated examples/train/search/searchr1_dataset.py --local_dir /mnt/local_storage/data/searchR1 # 2. Download the wiki-18 e5 index + corpus into the same dir, assemble e5_Flat.index, gunzip corpus. # 3. Start the local e5 retrieval server (separate faiss-gpu conda env) on :8000. -# Then launch: # export WANDB_API_KEY= -# bash examples/train/megatron/run_search_megatron_mimo_7b_specdecode.sh +# bash examples/train/spec_decode/run_search_megatron_mimo_7b_specdecode.sh MODEL_NAME="XiaomiMiMo/MiMo-7B-RL" # Dataset on the fast, non-persistent local disk (not the ~/default quota). diff --git a/skyrl/backends/skyrl_train/mtp/mtp_optim.py b/skyrl/backends/skyrl_train/mtp/mtp_optim.py index 5bd30c0413..43118fb60f 100644 --- a/skyrl/backends/skyrl_train/mtp/mtp_optim.py +++ b/skyrl/backends/skyrl_train/mtp/mtp_optim.py @@ -58,13 +58,13 @@ def freeze_mtp_params_pre_wrap(model_or_models: Union[torch.nn.Module, List[torc if is_mtp_param_name(name): param.requires_grad = False num_frozen += 1 - # Only registered when C-full is on, so zero matches means the `.mtp.` naming drifted -- the head - # would silently train in the policy grad buffer and defeat the isolation. Fail loud instead. + # Only registered when the separate MTP optimizer is on, so zero matches means the `.mtp.` naming + # drifted -- the head would silently train in the policy grad buffer and defeat the isolation. Fail loud. if num_frozen == 0: raise RuntimeError( "freeze_mtp_params_pre_wrap matched no parameters via is_mtp_param_name (expected '.mtp.' " - "or 'mtp.' in the name). The MTP submodule naming likely changed -- C-full grad isolation " - "would silently break. Update is_mtp_param_name in mtp/mtp_optim.py." + "or 'mtp.' in the name). The MTP submodule naming likely changed -- the head's grad " + "isolation would silently break. Update is_mtp_param_name in mtp/mtp_optim.py." ) return model_or_models @@ -112,7 +112,7 @@ def _build_mtp_ddp_config(ddp_config) -> DistributedDataParallelConfig: class MTPOptimizer: - """Owns the MTP/draft head's isolated grad buffer + ``DistributedOptimizer`` (C-full). + """Owns the MTP/draft head's isolated grad buffer + ``DistributedOptimizer``. The head stays inside the policy ``GPTModel`` (so capture / replay / weight-export are unchanged) but its params are excluded from the policy DDP buffer (via :func:`freeze_mtp_params_pre_wrap`) and wrapped diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index 5a15f378f8..ead7104601 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -379,15 +379,9 @@ def init_configs( provider = bridge.to_megatron_provider() - # Disable MTP for ordinary training: Megatron's in-forward MTP aux loss is unused, and under - # full recompute its checkpointed forward passes packed_seq_params positionally into - # tensor_parallel.checkpoint (tensors only), breaking packed-sequence backward. - # EXCEPTION: decoupled MTP/draft training (trainer.mtp.enabled on the policy worker) MUST keep - # the heads built — it runs them in-forward in eval mode (which skips the recompute path that - # breaks, see mtp/hidden_capture.py), scores a separate draft loss via a forward hook, and - # weight-syncs the heads to vLLM. It is compatible with sample packing (the draft loss masks - # cross-segment positions, see mtp/soft_ce.py::shift_mask_for_mtp), so we do NOT force packing - # off here. The MTP-config block below resolves the final head count. + # Disable MTP heads for ordinary training -- the native in-forward MTP loss breaks + # packed-sequence backward under full recompute. EXCEPTION: decoupled MTP/draft training + # (trainer.mtp.enabled) keeps them built. _mtp_cfg = getattr(self.cfg, "mtp", None) _mtp_training = enable_mtp and _mtp_cfg is not None and getattr(_mtp_cfg, "enabled", False) if not _mtp_training and getattr(provider, "mtp_num_layers", None): @@ -438,26 +432,15 @@ def init_configs( for k, v in transformer_config_kwargs.items(): setattr(provider, k, v) - # Multi-Token Prediction (MTP) head configuration. Model-agnostic: megatron-bridge populates - # provider.mtp_num_layers from whichever HF key the model uses — DeepSeek/GLM via - # num_nextn_predict_layers, Qwen3.5/Qwen3-Next via the generic - # (mtp_num_hidden_layers -> mtp_num_layers) config mapping. We layer SkyRL's explicit config - # on top: - # - enable_mtp=False (ref worker): force MTP off (the heads don't affect ref logprobs - # and would only waste compute/memory). - # - mtp_num_layers set: override the model default (0 force-disables). - # - mtp_num_layers=None: keep whatever the bridge inferred from the model. + # MTP head count: megatron-bridge infers provider.mtp_num_layers from the model's HF config. + # ref worker (enable_mtp=False) forces off; an explicit mtp_num_layers + # overrides (0 disables); None keeps the bridge's inferred count. if not enable_mtp: provider.mtp_num_layers = None elif megatron_config.mtp_num_layers is not None: provider.mtp_num_layers = megatron_config.mtp_num_layers or None - # When MTP training is requested (trainer.mtp.enabled), the model must actually resolve to - # >= 1 head — otherwise the draft loss would silently no-op (the wrapper keys off - # model_config.mtp_num_layers). _apply_mtp_config no longer forces a head count, so a model - # whose checkpoint ships no MTP layers lands here with mtp_num_layers=None; users who want - # to train fresh heads from scratch can still force-build them via - # policy.megatron_config.mtp_num_layers (note vLLM's `method: mtp` would also need the HF - # config to declare them). + # MTP training requires the model to resolve to >= 1 head, else the draft loss silently + # no-ops (the wrapper keys off model_config.mtp_num_layers) -- fail loud instead. mtp_cfg = getattr(self.cfg, "mtp", None) if ( enable_mtp @@ -472,18 +455,9 @@ def init_configs( "policy.megatron_config.mtp_num_layers to force-build fresh heads." ) if getattr(provider, "mtp_num_layers", None): - # The native MTP heads are still *built* (so the megatron-bridge weight mappings - # round-trip them to HF / vLLM), but SkyRL trains them with a decoupled, explicit - # loss (see MegatronModelWrapper) instead of Megatron's in-forward - # process_mtp_loss / MTPLossAutoScaler path. The explicit weight is - # megatron_config.mtp_loss_weight. - # - # CRITICAL: GPTModel/HybridModel.forward call process_mtp_loss unconditionally when MTP - # heads exist, and Megatron now DERIVES labels from input_ids ("e.g. RL training") so - # passing no labels no longer short-circuits it. Left active, that native hard-CE loss - # back-props into the policy trunk (inflated grad-norm, entropy collapse), independent of - # mtp_loss_weight. Patch process_mtp_loss to a no-op at its call sites so only SkyRL's - # decoupled loss trains the head. Must run before any forward (here, at config time). + # Heads stay built (for HF/vLLM weight round-trip) but SkyRL trains them with its own + # decoupled loss. Disable Megatron's native in-forward MTP loss (must run before any + # forward) or it back-props into the policy trunk and collapses entropy. See native_loss_patch.py. from skyrl.backends.skyrl_train.mtp.native_loss_patch import ( disable_native_mtp_loss, ) @@ -861,24 +835,21 @@ def init_model(self, model_path, num_training_steps: int = 1e9): logger.info("freeze_moe_router=True: freezing MoE router params") self.provider.register_pre_wrap_hook(freeze_moe_router) - # MTP C-full: isolate the draft head into its own grad buffer + optimizer so the policy's - # distributed grad reduction is byte-identical to a no-MTP model (the head's presence in the - # shared grad buffer otherwise perturbs the policy reduction and collapses entropy — see - # mtp/mtp_optim.py). Step 1, here: freeze the head BEFORE the policy DDP wrap so Megatron's - # DDP (which only buckets requires_grad params) excludes it from the policy grad buffer. The head - # is re-enabled and given its own buffer + optimizer after the policy optimizer is built. + # mtp_separate_optimizer: isolate the draft head into its own grad buffer + + # optimizer; see mtp/mtp_optim.py. Step 1: freeze the head before the policy DDP wrap so it's + # excluded from the policy grad buffer (Megatron only buckets requires_grad params). self._mtp_separate = None - self._mtp_cfull_enabled = bool( + self._mtp_separate_optim_enabled = bool( self.cfg.policy.megatron_config.mtp_separate_optimizer and getattr(self.provider, "mtp_num_layers", None) ) - if self._mtp_cfull_enabled: + if self._mtp_separate_optim_enabled: from skyrl.backends.skyrl_train.mtp.mtp_optim import ( freeze_mtp_params_pre_wrap, ) self.provider.register_pre_wrap_hook(freeze_mtp_params_pre_wrap) if self._rank == 0: - logger.info("MTP C-full: registered pre-wrap hook to isolate the draft head's grad buffer") + logger.info("MTP separate optimizer: registered pre-wrap hook to isolate the draft head's grad buffer") # wrap with DDP for training wrap_with_ddp = not self.cfg.policy.inference_only_init @@ -922,11 +893,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): num_training_steps=num_training_steps, ) - # MTP C-full step 2: re-enable the draft head and give it its OWN grad buffer + - # DistributedOptimizer (the policy optimizer above already excludes it — it was frozen - # before the policy wrap). The head co-trains at full strength while the policy reduction - # stays byte-identical to a no-MTP model. See mtp/mtp_optim.py. - if self._mtp_cfull_enabled: + # MTP separate optimizer step 2: re-enable the head and give it its own grad buffer + optimizer. + if self._mtp_separate_optim_enabled: from skyrl.backends.skyrl_train.mtp.mtp_optim import ( MTPOptimizer, ) @@ -944,7 +912,9 @@ def init_model(self, model_path, num_training_steps: int = 1e9): ) if self._rank == 0: n = sum(p.numel() for p in self._mtp_separate.mtp_params) - logger.info(f"MTP C-full: isolated draft head ({n:,} params) into its own grad buffer + optimizer") + logger.info( + f"MTP separate optimizer: isolated draft head ({n:,} params) into its own grad buffer + optimizer" + ) # create worker model self.model = MegatronModelWrapper( @@ -954,9 +924,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): policy_loss_fn=self.policy_loss_fn, ) - # MTP C-full step 3: hide the draft head from the POLICY finalize (which would otherwise coalesce - # the head's layernorm grads into the policy TP all-reduce, perturbing it). Must run AFTER the - # wrapper, which sets finalize_model_grads_func on this same model_config object. + # MTP separate optimizer step 3: hide the head from the policy finalize (must run after the wrapper sets + # finalize_model_grads_func). if self._mtp_separate is not None: from megatron.core.utils import get_model_config @@ -1108,7 +1077,7 @@ def forward_backward( for chunk in self.actor_module: # if use distributed optimizer, zero grad buffer will be handled by optimizer chunk.zero_grad_buffer() - # MTP C-full: zero the isolated draft-head grad buffer too (it lives outside self.actor_module). + # MTP separate optimizer: zero the isolated draft-head grad buffer too (it lives outside self.actor_module). if self._mtp_separate is not None: self._mtp_separate.zero_grad_buffer() @@ -1291,7 +1260,7 @@ def optim_step(self) -> Optional[float]: if self.optimizer is None: raise RuntimeError("optim_step called but policy.inference_only_init=True (no optimizer constructed)") - # MTP C-full: hide the draft head during the POLICY step so its grad-norm + clip EXCLUDE the + # MTP separate optimizer: hide the draft head during the POLICY step so its grad-norm + clip EXCLUDE the # head (Megatron's grad-norm iterates the GPTModel's requires_grad params and reads main_grad; # the head is structurally still in the GPTModel, so without this it dilutes the policy clip). # The policy step then matches a no-MTP run; the head is finalized + stepped separately after. @@ -1310,7 +1279,7 @@ def optim_step(self) -> Optional[float]: return grad_norm def save_checkpoint(self, ckpt_dir: str, tokenizer=None): - """Save the policy checkpoint, plus the C-full MTP head's separate optimizer state. + """Save the policy checkpoint, plus the MTP head's separate optimizer state. The isolated draft-head optimizer (``self._mtp_separate``) lives OUTSIDE ``self.optimizer``, so the strategy's checkpoint does not cover it. Persist its (DP-sharded) state per global rank @@ -1327,16 +1296,16 @@ def load_checkpoint(self, ckpt_dir: str, load_optimizer_states: bool = True, loa load_optimizer_states=load_optimizer_states, load_lr_scheduler_states=load_lr_scheduler_states, ) - # Restore the C-full MTP head's separate optimizer state. Guarded on existence so a checkpoint - # written without C-full (or a topology change) degrades to "fresh head optimizer" instead of - # crashing the resume. + # Restore the MTP head's separate optimizer state. Guarded on existence so a checkpoint + # written without the separate MTP optimizer (or a topology change) degrades to "fresh head + # optimizer" instead of crashing the resume. if self._mtp_separate is not None and load_optimizer_states: mtp_path = os.path.join(ckpt_dir, f"mtp_optim_rank{self._rank}.pt") if os.path.exists(mtp_path): self._mtp_separate.load_state_dict(torch.load(mtp_path, map_location="cpu")) elif self._rank == 0: logger.warning( - f"MTP C-full: no separate-optimizer checkpoint at {mtp_path}; " + f"MTP separate optimizer: no checkpoint at {mtp_path}; " "the draft head's optimizer state was not restored." ) return states diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 97e0345267..00f5c248ae 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -242,7 +242,7 @@ class MegatronConfig(BaseConfig): tied-embedding models (e.g. Qwen3.5), where a trainable shared head lets the draft loss nudge the policy's own logits. Set False to train the shared embedding/head with the draft loss too.""" mtp_separate_optimizer: bool = True - """C-full: give the MTP/draft head its OWN grad buffer + DistributedOptimizer, fully isolated from + """Give the MTP/draft head its OWN grad buffer + DistributedOptimizer, fully isolated from the policy's. The decoupled draft loss is autograd-clean, but when the head shares the policy's Megatron DDP grad buffer its mere presence changes the layout — and thus the floating-point result — of the policy gradient's distributed reduction (DP reduce-scatter + TP all-reduce of layernorm diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index d40f14a031..fc7ec28ffb 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -220,16 +220,10 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): # TODO (sumanthrh): Most of this should be moved to __post_init__ for the dataclasses def _apply_mtp_config(cfg: SkyRLTrainConfig): - """Propagate the high-level ``trainer.mtp`` knob to the training + inference configs. - - When MTP is enabled, train the model's native MTP heads (Megatron builds as many as the - checkpoint ships — typically one; override via ``policy.megatron_config.mtp_num_layers``) with - the decoupled draft loss, and enable vLLM MTP speculative decoding with - ``num_speculative_tokens`` draft tokens per step. The draft depth is deliberately decoupled - from the trained head count: vLLM drafts depth > 1 by reusing the (single) MTP head - autoregressively, so e.g. ``num_speculative_tokens=3`` with a 1-head checkpoint is valid and - does NOT build extra (randomly initialized) heads on the training side. When disabled, force - the MTP heads off so they are neither built nor run. + """Propagate the high-level ``trainer.mtp`` knob to the training + inference configs: train the + model's native MTP heads with the decoupled draft loss and enable vLLM MTP speculative decoding. + The vLLM draft depth (``num_speculative_tokens``) is decoupled from the trained head count + (depth > 1 reuses the head autoregressively). When disabled, force the heads off. """ mtp = getattr(cfg.trainer, "mtp", None) if mtp is None: @@ -247,16 +241,12 @@ def _apply_mtp_config(cfg: SkyRLTrainConfig): "trainer.mtp.enabled=true but trainer.policy.megatron_config.mtp_num_layers=0 " "(explicit force-disable). Remove the mtp_num_layers override or disable trainer.mtp." ) - # Training side: train the checkpoint's native MTP heads with the decoupled draft loss. - # mcfg.mtp_num_layers is intentionally left untouched (None => the megatron-bridge infers the - # head count from the model's own HF config); MegatronWorker fails loudly if the model resolves - # to zero heads while MTP is enabled. + # Leave mcfg.mtp_num_layers untouched (None => megatron-bridge infers the head count from the + # model's HF config; MegatronWorker fails loud if it resolves to zero while MTP is enabled). mcfg.mtp_loss_type = mtp.loss_type mcfg.mtp_loss_weight = mtp.loss_weight - # DEBUG GUARD: SKYRL_DISABLE_SPEC=1 keeps MTP head TRAINING on (the decoupled draft loss set - # above) but leaves vLLM rollout in plain autoregressive (lossless) decode — isolates the - # MTP-training machinery from the spec-decode rollout sampler. + # SKYRL_DISABLE_SPEC=1: keep MTP head training on but leave vLLM rollout plain autoregressive. if os.environ.get("SKYRL_DISABLE_SPEC") == "1": return