diff --git a/docs/assets/rl-kernel/rocm-performance-trajectories.png b/docs/assets/rl-kernel/rocm-performance-trajectories.png new file mode 100644 index 000000000..c258327a0 Binary files /dev/null and b/docs/assets/rl-kernel/rocm-performance-trajectories.png differ diff --git a/docs/assets/rl-kernel/rocm-training-stability.png b/docs/assets/rl-kernel/rocm-training-stability.png new file mode 100644 index 000000000..6883983ee Binary files /dev/null and b/docs/assets/rl-kernel/rocm-training-stability.png differ diff --git a/tests/test_linear_logp_provider.py b/tests/test_linear_logp_provider.py new file mode 100644 index 000000000..ddc82a0ed --- /dev/null +++ b/tests/test_linear_logp_provider.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace + +import pytest +import torch + +from vime.backends.megatron_utils.linear_logp_provider import ( + LinearLogpContext, + LinearLogpProviderUnavailable, + LinearLogpRequest, + LinearLogpResult, + LinearProjection, + TokenLayout, + VocabPartition, + compute_linear_logp, +) + + +def _context() -> LinearLogpContext: + return LinearLogpContext( + hidden=torch.randn(3, 4), + projection=LinearProjection(weight=torch.randn(5, 4), bias=torch.randn(5)), + vocab_partition=VocabPartition(local_start=0, local_size=5, real_size=4, padded_size=5), + ) + + +def _request(*, with_entropy=False, context=None) -> LinearLogpRequest: + logits = torch.randn(3, 5, requires_grad=True) + return LinearLogpRequest( + logits=logits, + target_ids=torch.tensor([1, 2, 3]), + tensor_parallel_group=None, + token_layout=TokenLayout(world_size=2, rank=1, layout="zigzag"), + with_entropy=with_entropy, + with_entropy_grad=with_entropy, + chunk_size=64, + context=context, + ) + + +def _native(logits, *_args, with_entropy, **_kwargs): + return logits[:, :1], logits.sum(dim=-1) if with_entropy else None + + +def _args(path=None, mode="auto"): + return SimpleNamespace(linear_logp_provider=path, linear_logp_provider_mode=mode) + + +def _install_provider(monkeypatch, provider): + module = types.ModuleType("linear_logp_provider_fixture") + module.provider = provider + monkeypatch.setitem(sys.modules, module.__name__, module) + return f"{module.__name__}.provider" + + +@pytest.mark.unit +def test_unconfigured_provider_uses_native_path(): + request = _request() + + logp, entropy = compute_linear_logp(args=_args(), request=request, native=_native) + + torch.testing.assert_close(logp, request.logits[:, :1]) + assert entropy is None + + +@pytest.mark.unit +def test_provider_receives_structural_context(monkeypatch): + context = _context() + request = _request(with_entropy=True, context=context) + observed = {} + + def provider(actual_request): + observed["context"] = actual_request.context + return LinearLogpResult( + logp=actual_request.logits[:, :1], + entropy=actual_request.logits.sum(dim=-1), + backend_id="fixture", + contract_id="fixture.v1", + ) + + path = _install_provider(monkeypatch, provider) + logp, entropy = compute_linear_logp(args=_args(path, "strict"), request=request, native=_native) + + assert observed["context"] is context + torch.testing.assert_close(logp, request.logits[:, :1]) + torch.testing.assert_close(entropy, request.logits.sum(dim=-1)) + + +@pytest.mark.unit +def test_auto_falls_back_only_for_explicit_unavailability(monkeypatch): + request = _request() + + def unavailable(_request): + raise LinearLogpProviderUnavailable("unsupported") + + path = _install_provider(monkeypatch, unavailable) + logp, _ = compute_linear_logp(args=_args(path), request=request, native=_native) + torch.testing.assert_close(logp, request.logits[:, :1]) + + def broken(_request): + raise RuntimeError("provider bug") + + path = _install_provider(monkeypatch, broken) + with pytest.raises(RuntimeError, match="provider bug"): + compute_linear_logp(args=_args(path), request=request, native=_native) + + +@pytest.mark.unit +def test_strict_mode_fails_when_provider_is_unavailable(monkeypatch): + def provider(_request): + raise LinearLogpProviderUnavailable("unsupported") + + path = _install_provider(monkeypatch, provider) + with pytest.raises(RuntimeError, match="is unavailable"): + compute_linear_logp(args=_args(path, "strict"), request=_request(), native=_native) + + +@pytest.mark.unit +def test_strict_mode_validates_identity_and_autograd(monkeypatch): + request = _request() + + def provider(actual_request): + return { + "logp": actual_request.logits[:, :1].detach(), + "entropy": None, + "backend_id": "fixture", + "contract_id": "fixture.v1", + "provenance": {}, + } + + path = _install_provider(monkeypatch, provider) + with pytest.raises(ValueError, match="detached"): + compute_linear_logp(args=_args(path, "strict"), request=request, native=_native) + + +@pytest.mark.unit +def test_structural_context_rejects_misaligned_projection(): + with pytest.raises(ValueError, match="projection and vocabulary"): + LinearLogpContext( + hidden=torch.randn(3, 4), + projection=LinearProjection(weight=torch.randn(5, 4)), + vocab_partition=VocabPartition(local_start=0, local_size=4, real_size=5, padded_size=5), + ) diff --git a/tests/test_logprob_response_spans.py b/tests/test_logprob_response_spans.py index d22ef4591..8492b6d68 100644 --- a/tests/test_logprob_response_spans.py +++ b/tests/test_logprob_response_spans.py @@ -3,10 +3,14 @@ import _cp_dist_helpers # noqa: F401 import pytest import torch - from megatron.core import mpu -from vime.backends.megatron_utils.loss import _build_topp_keep_mask, get_rollout_top_p_logprob_kwargs +from vime.backends.megatron_utils import loss as loss_module +from vime.backends.megatron_utils.loss import ( + _build_topp_keep_mask, + get_log_probs_and_entropy, + get_rollout_top_p_logprob_kwargs, +) NUM_GPUS = 0 @@ -95,5 +99,52 @@ def test_top_p_mask_aligns_with_cp1_response_rows(monkeypatch): assert masked_rows == {2: [13], 3: [14], 5: [21], 6: [22], 7: [23]} +@pytest.mark.unit +def test_provider_receives_unscaled_logits_and_native_fallback_scales_once(monkeypatch): + _set_cp(monkeypatch, size=1, rank=0) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_group", lambda: None, raising=False) + monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: 1, raising=False) + observed = {} + + def calculate(logits, *_args, with_entropy, **_kwargs): + observed["native_logits"] = logits.clone() + return logits[:, :1], logits.sum(dim=-1) if with_entropy else None + + def dispatch(*, request, native, **_kwargs): + observed["request_logits"] = request.logits.clone() + observed["temperature"] = request.temperature + return native( + request.logits, + request.target_ids, + request.tensor_parallel_group, + with_entropy=request.with_entropy, + with_entropy_grad=request.with_entropy_grad, + chunk_size=request.chunk_size, + log_prob_keep_mask=request.log_prob_keep_mask, + ) + + monkeypatch.setattr(loss_module, "calculate_log_probs_and_entropy", calculate) + monkeypatch.setattr(loss_module, "compute_linear_logp", dispatch) + logits = torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) + args = Namespace( + allgather_cp=False, + entropy_coef=0.0, + log_probs_chunk_size=-1, + rollout_temperature=0.5, + ) + + get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=[torch.tensor([0, 1, 2])], + total_lengths=[3], + response_lengths=[2], + ) + + torch.testing.assert_close(observed["request_logits"], logits.squeeze(0)) + torch.testing.assert_close(observed["native_logits"], logits.squeeze(0) / 0.5) + assert observed["temperature"] == 0.5 + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index 05b6e7b2d..22d889691 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -1,3 +1,4 @@ +import argparse import importlib.util import sys import types @@ -167,6 +168,37 @@ def test_allgather_cp_ignores_cp_size_one(monkeypatch): module._validate_allgather_cp_supported(args) +@pytest.mark.unit +def test_strict_linear_logp_provider_requires_provider_path(monkeypatch): + module = load_arguments_module(monkeypatch) + args = types.SimpleNamespace(linear_logp_provider=None, linear_logp_provider_mode="strict") + + with pytest.raises(ValueError, match="requires --linear-logp-provider"): + module._validate_linear_logp_provider_args(args) + + +@pytest.mark.unit +def test_linear_logp_provider_arguments_are_registered(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + module.RouterArgs = types.SimpleNamespace(add_cli_args=lambda parser, **_kwargs: parser) + parser = argparse.ArgumentParser() + module.get_vime_extra_args_provider()(parser) + + args = parser.parse_args( + [ + "--rollout-batch-size", + "1", + "--linear-logp-provider", + "package.module.provider", + "--linear-logp-provider-mode", + "strict", + ] + ) + + assert args.linear_logp_provider == "package.module.provider" + assert args.linear_logp_provider_mode == "strict" + + @pytest.mark.unit def test_update_weight_disk_dir_required_for_disk_transport(monkeypatch): module = load_vime_arguments_module(monkeypatch) diff --git a/vime/backends/megatron_utils/arguments.py b/vime/backends/megatron_utils/arguments.py index 5176eb8fc..9653b958e 100644 --- a/vime/backends/megatron_utils/arguments.py +++ b/vime/backends/megatron_utils/arguments.py @@ -73,10 +73,18 @@ def _is_moe_config(hf_config): ) +def _validate_linear_logp_provider_args(args): + provider = str(getattr(args, "linear_logp_provider", "") or "").strip() + mode = str(getattr(args, "linear_logp_provider_mode", "auto")).strip().lower() + if mode == "strict" and not provider: + raise ValueError("--linear-logp-provider-mode strict requires --linear-logp-provider") + + def validate_args(args): """Run megatron's own validate_args plus vime-specific megatron validations.""" _megatron_validate_args(args) + _validate_linear_logp_provider_args(args) # always use varlen args.variable_seq_lengths = True diff --git a/vime/backends/megatron_utils/linear_logp_provider.py b/vime/backends/megatron_utils/linear_logp_provider.py new file mode 100644 index 000000000..35e8c828a --- /dev/null +++ b/vime/backends/megatron_utils/linear_logp_provider.py @@ -0,0 +1,242 @@ +"""Optional provider boundary for Megatron ``linear_logp`` computation. + +Vime owns token layout and loss composition. A provider owns the numerical +implementation, tensor-parallel reduction, and runtime provenance. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Any, Literal + +import torch + +ProviderMode = Literal["auto", "strict"] +TokenLayoutKind = Literal["single", "zigzag", "allgather"] + + +class LinearLogpProviderUnavailable(RuntimeError): + """Ask Vime to use its native implementation when mode is ``auto``.""" + + linear_logp_provider_unavailable = True + + +@dataclass(frozen=True) +class TokenLayout: + world_size: int + rank: int + layout: TokenLayoutKind + + def __post_init__(self) -> None: + if self.world_size < 1 or not 0 <= self.rank < self.world_size: + raise ValueError(f"invalid token topology: world_size={self.world_size}, rank={self.rank}") + if (self.world_size == 1) != (self.layout == "single"): + raise ValueError("single-rank token layout must be 'single'; multi-rank layout must not be 'single'") + + +@dataclass(frozen=True) +class LinearProjection: + weight: torch.Tensor + bias: torch.Tensor | None = None + + def __post_init__(self) -> None: + if not isinstance(self.weight, torch.Tensor) or self.weight.ndim != 2: + raise ValueError("linear_logp projection weight must be rank 2") + if self.bias is not None and ( + not isinstance(self.bias, torch.Tensor) + or self.bias.shape != (self.weight.size(0),) + or self.bias.device != self.weight.device + ): + raise ValueError("linear_logp projection bias must match the local vocabulary shard") + + +@dataclass(frozen=True) +class VocabPartition: + local_start: int + local_size: int + real_size: int + padded_size: int + + def __post_init__(self) -> None: + if self.local_start < 0 or self.local_size <= 0: + raise ValueError("linear_logp vocabulary partition is empty or negative") + if self.real_size <= 0 or self.real_size > self.padded_size: + raise ValueError("linear_logp global vocabulary sizes are invalid") + if self.local_start + self.local_size > self.padded_size: + raise ValueError("linear_logp local vocabulary shard exceeds the padded vocabulary") + + +@dataclass(frozen=True) +class LinearLogpContext: + hidden: torch.Tensor + projection: LinearProjection + vocab_partition: VocabPartition + + def __post_init__(self) -> None: + if not isinstance(self.hidden, torch.Tensor) or self.hidden.ndim != 2: + raise ValueError("linear_logp hidden states must have shape [T, H]") + if self.hidden.size(1) != self.projection.weight.size(1): + raise ValueError("linear_logp hidden and projection widths do not match") + if self.projection.weight.size(0) != self.vocab_partition.local_size: + raise ValueError("linear_logp projection and vocabulary shard widths do not match") + if self.hidden.device != self.projection.weight.device: + raise ValueError("linear_logp context tensors must share a device") + + +@dataclass(frozen=True) +class LinearLogpRequest: + logits: torch.Tensor + target_ids: torch.Tensor + tensor_parallel_group: Any + token_layout: TokenLayout + with_entropy: bool + with_entropy_grad: bool + chunk_size: int + log_prob_keep_mask: torch.Tensor | None = None + context: LinearLogpContext | None = None + temperature: float | torch.Tensor | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.logits, torch.Tensor) or self.logits.ndim != 2: + raise ValueError("linear_logp logits must have shape [T, V_local]") + if self.target_ids.shape != (self.logits.size(0),): + raise ValueError("linear_logp target IDs must have shape [T]") + if self.target_ids.device != self.logits.device: + raise ValueError("linear_logp target IDs and logits must share a device") + if self.target_ids.is_floating_point() or self.target_ids.is_complex(): + raise TypeError("linear_logp target IDs must use an integer dtype") + if self.log_prob_keep_mask is not None and self.log_prob_keep_mask.shape != self.logits.shape: + raise ValueError("linear_logp keep mask must match logits") + if self.context is not None and ( + self.context.hidden.size(0) != self.logits.size(0) + or self.context.projection.weight.size(0) != self.logits.size(1) + or self.context.hidden.device != self.logits.device + ): + raise ValueError("linear_logp structural context does not match logits") + + +@dataclass(frozen=True) +class LinearLogpResult: + logp: torch.Tensor + entropy: torch.Tensor | None + backend_id: str + contract_id: str + provenance: Mapping[str, Any] = field(default_factory=dict) + + +def linear_logp_provider_path(args: Any) -> str | None: + path = str(getattr(args, "linear_logp_provider", "") or "").strip() + return path or None + + +def linear_logp_provider_mode(args: Any) -> ProviderMode: + mode = str(getattr(args, "linear_logp_provider_mode", "auto")).strip().lower() + if mode not in {"auto", "strict"}: + raise ValueError(f"linear_logp provider mode must be 'auto' or 'strict', got {mode!r}") + return mode # type: ignore[return-value] + + +def compute_linear_logp( + *, + args: Any, + request: LinearLogpRequest, + native: Callable[..., tuple[torch.Tensor, torch.Tensor | None]], +) -> tuple[torch.Tensor, torch.Tensor | None]: + path = linear_logp_provider_path(args) + if path is None: + return _native(request, native) + + mode = linear_logp_provider_mode(args) + try: + provider = _load_provider(path) + result = provider(request) + except Exception as exc: + unavailable = isinstance(exc, (ImportError, AttributeError)) or bool( + getattr(exc, "linear_logp_provider_unavailable", False) + ) + if not unavailable or mode == "strict": + if unavailable and mode == "strict": + raise RuntimeError(f"linear_logp provider {path!r} is unavailable: {exc}") from exc + raise + return _native(request, native) + + normalized = _normalize_result(result) + _validate_result(normalized, request, strict=mode == "strict") + return normalized.logp, normalized.entropy + + +def _load_provider(path: str) -> Callable[[LinearLogpRequest], Any]: + module_path, separator, attribute = path.rpartition(".") + if not separator: + raise ImportError("linear_logp provider must be a fully qualified import path") + provider = getattr(importlib.import_module(module_path), attribute) + if not callable(provider): + raise TypeError(f"linear_logp provider {path!r} is not callable") + return provider + + +def _native( + request: LinearLogpRequest, + native: Callable[..., tuple[torch.Tensor, torch.Tensor | None]], +) -> tuple[torch.Tensor, torch.Tensor | None]: + return native( + request.logits, + request.target_ids, + request.tensor_parallel_group, + with_entropy=request.with_entropy, + with_entropy_grad=request.with_entropy_grad, + chunk_size=request.chunk_size, + log_prob_keep_mask=request.log_prob_keep_mask, + ) + + +def _normalize_result(result: Any) -> LinearLogpResult: + if isinstance(result, LinearLogpResult): + return result + getter = result.__getitem__ if isinstance(result, Mapping) else lambda name: getattr(result, name) + try: + return LinearLogpResult( + logp=getter("logp"), + entropy=getter("entropy"), + backend_id=getter("backend_id"), + contract_id=getter("contract_id"), + provenance=getter("provenance"), + ) + except (AttributeError, KeyError) as exc: + raise TypeError("linear_logp provider returned an invalid result") from exc + + +def _validate_result(result: LinearLogpResult, request: LinearLogpRequest, *, strict: bool) -> None: + if result.logp.shape != (request.logits.size(0), 1): + raise ValueError("linear_logp provider must return logp with shape [T, 1]") + if result.logp.device != request.logits.device or not result.logp.is_floating_point(): + raise ValueError("linear_logp provider returned logp with an invalid device or dtype") + if request.with_entropy: + if result.entropy is None or result.entropy.shape != (request.logits.size(0),): + raise ValueError("linear_logp provider must return entropy with shape [T]") + elif result.entropy is not None: + raise ValueError("linear_logp provider returned entropy when it was not requested") + if strict: + if not result.backend_id or not result.contract_id: + raise ValueError("strict linear_logp results require backend and contract IDs") + if request.logits.requires_grad and not result.logp.requires_grad: + raise ValueError("strict linear_logp result is detached from autograd") + if request.with_entropy_grad and result.entropy is not None and not result.entropy.requires_grad: + raise ValueError("strict linear_logp entropy is detached from autograd") + + +__all__ = [ + "LinearLogpContext", + "LinearLogpProviderUnavailable", + "LinearLogpRequest", + "LinearLogpResult", + "LinearProjection", + "TokenLayout", + "VocabPartition", + "compute_linear_logp", + "linear_logp_provider_mode", + "linear_logp_provider_path", +] diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 566be1df7..9da9f78bd 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -30,6 +30,7 @@ get_sum_of_sample_mean, slice_log_prob_with_cp, ) +from .linear_logp_provider import LinearLogpContext, LinearLogpRequest, TokenLayout, compute_linear_logp ROLLOUT_TOP_P_TOKEN_KEYS = ( "rollout_top_p_token_ids", @@ -521,6 +522,7 @@ def get_log_probs_and_entropy( non_loss_data: bool = True, top_p_token_ids: list[list[int]] | None = None, top_p_token_offsets: list[list[int]] | None = None, + linear_logp_context: LinearLogpContext | None = None, ) -> dict[str, list[torch.Tensor]]: """Compute per-token log-probabilities (and optionally entropy) on responses. @@ -532,15 +534,13 @@ def get_log_probs_and_entropy( log-probabilities; entropy is always computed from the unmasked logits. """ assert non_loss_data - assert logits.dtype == torch.float32, f"{logits.dtype}" + if logits.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError(f"linear_logp logits must be floating point, got {logits.dtype}") assert len(logits.shape) == 3, f"{logits.shape}" assert logits.size(0) == 1, f"{logits.shape}" logits = logits.squeeze(0) - # Apply rollout temperature scaling to logits to match rollout-time log-probs. rollout_temperature = getattr(args, "rollout_temperature", 1.0) - if rollout_temperature != 1.0: - logits = logits / rollout_temperature logits = logits.contiguous() T = logits.size(0) device = logits.device @@ -567,15 +567,39 @@ def get_log_probs_and_entropy( args.allgather_cp, ) - # --- compute on full [T,V] logits at once via calculate_log_probs_and_entropy --- - log_prob_full, entropy_full = calculate_log_probs_and_entropy( - logits, - full_tokens, - tp_group, + cp_world_size = mpu.get_context_parallel_world_size() + request = LinearLogpRequest( + logits=logits, + target_ids=full_tokens, + tensor_parallel_group=tp_group, + token_layout=TokenLayout( + world_size=cp_world_size, + rank=mpu.get_context_parallel_rank(), + layout="single" if cp_world_size == 1 else "allgather" if args.allgather_cp else "zigzag", + ), with_entropy=with_entropy, with_entropy_grad=with_entropy_grad, chunk_size=chunk_size, log_prob_keep_mask=top_p_keep_mask, + context=linear_logp_context, + temperature=rollout_temperature, + metadata={ + "real_vocab_size": getattr(args, "vocab_size", None), + "padded_vocab_size": getattr(args, "padded_vocab_size", None), + "tp_rank": mpu.get_tensor_model_parallel_rank(), + "tp_world_size": mpu.get_tensor_model_parallel_world_size(), + }, + ) + + def native_linear_logp(native_logits, *native_args, **native_kwargs): + if rollout_temperature != 1.0: + native_logits = native_logits / rollout_temperature + return calculate_log_probs_and_entropy(native_logits, *native_args, **native_kwargs) + + log_prob_full, entropy_full = compute_linear_logp( + args=args, + request=request, + native=native_linear_logp, ) log_prob_full = log_prob_full.squeeze(-1) # [T, 1] -> [T] @@ -936,6 +960,7 @@ def policy_loss_function( batch: RolloutBatch, logits: torch.Tensor, sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], + linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute policy loss (PPO/GSPO) and metrics. @@ -974,6 +999,7 @@ def policy_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=True, + linear_logp_context=linear_logp_context, **get_rollout_top_p_logprob_kwargs(args, batch), ) @@ -1235,6 +1261,7 @@ def sft_loss_function( batch: RolloutBatch, logits: torch.Tensor, sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], + linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Compute supervised fine-tuning loss over response tokens. @@ -1262,6 +1289,7 @@ def sft_loss_function( total_lengths=total_lengths, response_lengths=response_lengths, with_entropy=False, + linear_logp_context=linear_logp_context, ) log_probs = log_probs_and_entropy["log_probs"] @@ -1286,6 +1314,7 @@ def loss_function( num_microbatches: int, step_global_batch_size: int, logits: torch.Tensor, + linear_logp_context: LinearLogpContext | None = None, ) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]: """Dispatch to the configured loss and rescale for Megatron integration. @@ -1336,10 +1365,21 @@ def loss_function( case _: raise ValueError(f"Unknown loss type: {args.loss_type}") + provider_kwargs = ( + {"linear_logp_context": linear_logp_context} if args.loss_type in {"policy_loss", "sft_loss"} else {} + ) if args.recompute_loss_function: - loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean, use_reentrant=False) + loss, log = checkpoint( + func, + args, + batch, + logits, + sum_of_sample_mean, + **provider_kwargs, + use_reentrant=False, + ) else: - loss, log = func(args, batch, logits, sum_of_sample_mean) + loss, log = func(args, batch, logits, sum_of_sample_mean, **provider_kwargs) # With allgather-CP, some CP ranks may have no loss-contributing tokens (e.g., all # padding). Without this, gradient doesn't flow through their attention path, so diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index 15ce22cff..2a1a6f19a 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -28,12 +28,14 @@ from megatron.core.pipeline_parallel.utils import unwrap_model except ImportError: from megatron.core.utils import unwrap_model + from vime.utils import logging_utils from vime.utils.memory_utils import clear_memory from .checkpoint import load_checkpoint, save_checkpoint from .cp_utils import reduce_train_step_metrics from .data import DataIterator, get_batch +from .linear_logp_provider import LinearLogpContext, LinearProjection, VocabPartition, linear_logp_provider_path from .loss import ROLLOUT_TOP_P_TOKEN_KEYS, get_rollout_top_p_logprob_kwargs, loss_function from .model_provider import get_model_provider_func from .stateless_adam import StatelessAdam @@ -41,6 +43,67 @@ logger = logging.getLogger(__name__) +def _unwrap_model_chunk(model_chunk): + while hasattr(model_chunk, "module"): + model_chunk = model_chunk.module + return model_chunk + + +def _install_linear_logp_capture(model_chunks: Sequence[DDP], args: Namespace) -> None: + from megatron.core.tensor_parallel.utils import VocabUtility + + for model_chunk in model_chunks: + owner = _unwrap_model_chunk(model_chunk) + output_layer = getattr(owner, "output_layer", None) + if output_layer is None or hasattr(output_layer, "_vime_linear_logp_capture_handle"): + continue + if not hasattr(output_layer, "output_size_per_partition") or not hasattr(output_layer, "tp_group"): + continue + + def capture(module, inputs, kwargs, *, context_owner=owner): + hidden = inputs[0] if inputs else kwargs.get("input_") + weight = inputs[1] if len(inputs) > 1 else kwargs.get("weight") + if weight is None: + weight = getattr(module, "weight", None) + if not isinstance(hidden, torch.Tensor) or not isinstance(weight, torch.Tensor): + raise RuntimeError("Megatron output layer did not expose hidden states and weight") + if hidden.ndim == 3: + hidden = hidden.transpose(0, 1).contiguous().reshape(-1, hidden.size(-1)) + elif hidden.ndim != 2: + raise RuntimeError(f"unsupported Megatron output-layer input shape: {tuple(hidden.shape)}") + + tp_world_size = mpu.get_tensor_model_parallel_world_size() + tp_rank = mpu.get_tensor_model_parallel_rank() + padded_vocab_size = int(getattr(args, "padded_vocab_size", weight.size(0) * tp_world_size)) + real_vocab_size = int(getattr(args, "vocab_size", padded_vocab_size)) + if weight.size(0) * tp_world_size != padded_vocab_size: + raise RuntimeError("linear_logp provider requires equal contiguous TP vocabulary shards") + vocab_start, _ = VocabUtility.vocab_range_from_per_partition_vocab_size( + weight.size(0), tp_rank, tp_world_size + ) + context_owner._vime_linear_logp_context = LinearLogpContext( + hidden=hidden, + projection=LinearProjection(weight=weight, bias=getattr(module, "bias", None)), + vocab_partition=VocabPartition( + local_start=int(vocab_start), + local_size=weight.size(0), + real_size=real_vocab_size, + padded_size=padded_vocab_size, + ), + ) + + output_layer._vime_linear_logp_capture_handle = output_layer.register_forward_pre_hook( + capture, with_kwargs=True + ) + + +def _take_linear_logp_context(model_chunk): + owner = _unwrap_model_chunk(model_chunk) + context = getattr(owner, "_vime_linear_logp_context", None) + owner._vime_linear_logp_context = None + return context + + def _disable_tqdm_for_non_main_rank() -> bool: return not ( mpu.get_data_parallel_rank(with_context_parallel=True) == 0 @@ -292,6 +355,8 @@ def setup_model_and_optimizer( assert args.load is not None or args.pretrained_checkpoint is not None model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder) + if role == "actor" and linear_logp_provider_path(args) is not None: + _install_linear_logp_capture(model, args) # Optimizer kwargs = {} @@ -431,6 +496,7 @@ def forward_step( if batch["multimodal_train_inputs"] is not None: forward_kwargs.update(batch["multimodal_train_inputs"]) output_tensor = model(**forward_kwargs) + linear_logp_context = _take_linear_logp_context(model) output_kwargs = { "args": args, @@ -438,6 +504,7 @@ def forward_step( "total_lengths": total_lengths, "response_lengths": response_lengths, "with_entropy": args.use_rollout_entropy, + "linear_logp_context": linear_logp_context, } if use_rollout_top_p_replay: output_kwargs.update(get_rollout_top_p_logprob_kwargs(args, batch)) @@ -650,10 +717,19 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p output_tensor = model(**forward_kwargs) + linear_logp_context = _take_linear_logp_context(model) + if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": os.environ["ROUTING_REPLAY_STAGE"] = old_stage - return output_tensor, partial(loss_function, args, batch, num_microbatches, step_global_batch_size) + return output_tensor, partial( + loss_function, + args, + batch, + num_microbatches, + step_global_batch_size, + linear_logp_context=linear_logp_context, + ) # Forward pass. forward_backward_func = get_forward_backward_func() diff --git a/vime/backends/vllm_utils/arguments.py b/vime/backends/vllm_utils/arguments.py index 7e7123a42..9d62ded74 100644 --- a/vime/backends/vllm_utils/arguments.py +++ b/vime/backends/vllm_utils/arguments.py @@ -1,7 +1,12 @@ import argparse from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.utils.argparse_utils import FlexibleArgumentParser + +try: + from vllm.utils.argparse_utils import FlexibleArgumentParser +except ImportError: + from vllm.utils import FlexibleArgumentParser + from vllm_router.launch_router import RouterArgs from vime.utils.http_utils import _wrap_ipv6 @@ -83,6 +88,7 @@ def wrapper(*name_or_flags, **kwargs): new_flags.append(s) new_kwargs = kwargs.copy() + new_kwargs.pop("deprecated", None) if "dest" in new_kwargs and isinstance(new_kwargs["dest"], str): if not new_kwargs["dest"].startswith("vllm_"): new_kwargs["dest"] = f"vllm_{new_kwargs['dest']}" @@ -98,10 +104,15 @@ def patched_add_argument_group(*g_args, **g_kwargs): parser.add_argument = _wrap_add_argument(old_add_argument) parser.add_argument_group = patched_add_argument_group - AsyncEngineArgs.add_cli_args(parser) - from vllm.entrypoints.openai.cli_args import FrontendArgs - - FrontendArgs.add_cli_args(parser) + try: + from vllm.entrypoints.launchers.cli_args import FrontendArgs + except ImportError: + from vllm.entrypoints.openai.cli_args import make_arg_parser + + make_arg_parser(parser) + else: + AsyncEngineArgs.add_cli_args(parser) + FrontendArgs.add_cli_args(parser) parser.add_argument = old_add_argument parser.add_argument_group = old_add_argument_group diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 423673bd1..9d8d53afa 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -93,7 +93,11 @@ def _run_vllm_server(kwargs: dict, env: dict) -> None: os.environ.update(env) from vllm.entrypoints.cli.serve import ServeSubcommand - from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args + + try: + from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args + except ImportError: + from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.utils.argparse_utils import FlexibleArgumentParser ns = argparse.Namespace(**kwargs) @@ -747,7 +751,15 @@ def _compute_server_args( def _vllm_server_field_names() -> frozenset[str]: """Return the vLLM fields accepted by CLI generation and config overrides.""" from vllm.engine.arg_utils import AsyncEngineArgs - from vllm.entrypoints.openai.cli_args import FrontendArgs + + try: + from vllm.entrypoints.launchers.cli_args import FrontendArgs + except ImportError: + from vllm.entrypoints.openai.cli_args import make_arg_parser + from vllm.utils.argparse_utils import FlexibleArgumentParser + + parser = make_arg_parser(FlexibleArgumentParser(add_help=False)) + return frozenset(action.dest for action in parser._actions if action.dest != argparse.SUPPRESS) return frozenset(f.name for f in (*dataclasses.fields(AsyncEngineArgs), *dataclasses.fields(FrontendArgs))) diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 113ab12a0..17d38be07 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -253,6 +253,18 @@ def add_train_arguments(parser): parser.add_argument( "--log-probs-chunk-size", type=int, default=-1, help="Chunk size to compute log probs to save memory" ) + parser.add_argument( + "--linear-logp-provider", + type=str, + default=None, + help="Fully qualified callable that replaces Megatron linear-logp computation.", + ) + parser.add_argument( + "--linear-logp-provider-mode", + choices=("auto", "strict"), + default="auto", + help="Fall back when the provider is unavailable, or fail closed in strict mode.", + ) parser.add_argument( "--only-train-params-name-list", type=str,