From e1eb5e276dbde11d6305122794466cc30fca3a7a Mon Sep 17 00:00:00 2001 From: "tianle.zhong" Date: Mon, 3 Aug 2026 23:37:47 +0000 Subject: [PATCH 1/3] [ops] feat: Support GPT-OSS learnable attention sinks --- .../test_flash_attention.py | 80 ++++++++++++++++++- tests/test_model_config.py | 34 ++++++++ vexact/batch_invariant_ops/README.md | 2 +- vexact/batch_invariant_ops/flash_attention.py | 15 ++++ vexact/config.py | 5 ++ 5 files changed, 134 insertions(+), 2 deletions(-) create mode 100644 tests/test_model_config.py diff --git a/tests/batch_invariant_ops/test_flash_attention.py b/tests/batch_invariant_ops/test_flash_attention.py index c926e26..20be600 100644 --- a/tests/batch_invariant_ops/test_flash_attention.py +++ b/tests/batch_invariant_ops/test_flash_attention.py @@ -15,7 +15,11 @@ import pytest import torch -from vexact.batch_invariant_ops.flash_attention import _get_fa_window_size, flash_attention_forward +from vexact.batch_invariant_ops.flash_attention import ( + _get_fa_learnable_sink, + _get_fa_window_size, + flash_attention_forward, +) from vexact.batch_invariant_ops.kv_cache_context import set_kv_cache_context @@ -37,6 +41,21 @@ def test_get_fa_window_size_rejects_non_positive_values(sliding_window): _get_fa_window_size(sliding_window) +def test_get_fa_learnable_sink_casts_to_query_dtype(): + query = torch.empty(1, dtype=torch.bfloat16) + sink = torch.randn(2, dtype=torch.float32) + + actual = _get_fa_learnable_sink(sink, query, use_cute=True) + + assert actual.dtype == query.dtype + torch.testing.assert_close(actual.float(), sink.to(torch.bfloat16).float()) + + +def test_get_fa_learnable_sink_rejects_fa3(): + with pytest.raises(NotImplementedError, match="require the FA4"): + _get_fa_learnable_sink(torch.zeros(2), torch.empty(1), use_cute=False) + + @pytest.mark.parametrize("use_cute", [False, True], ids=["fa3", "fa4"]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_flash_attention_sliding_window_matches_reference(use_cute): @@ -97,3 +116,62 @@ def test_flash_attention_sliding_window_matches_reference(use_cute): expected = expected.transpose(1, 2).reshape_as(actual) torch.testing.assert_close(actual.float(), expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_fa4_learnable_sink_matches_reference(): + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("FA4 requires SM90+") + pytest.importorskip("flash_attn.cute") + + torch.manual_seed(1) + device = torch.device("cuda") + dtype = torch.bfloat16 + batch_size, seq_len, num_heads, head_dim = 1, 16, 2, 64 + page_size = 256 + sliding_window = 4 + scaling = head_dim**-0.5 + + query = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) + key = torch.randn_like(query) + value = torch.randn_like(query) + sinks = torch.randn(num_heads, device=device, dtype=dtype) + key_cache = torch.zeros(1, page_size, num_heads, head_dim, device=device, dtype=dtype) + value_cache = torch.zeros_like(key_cache) + set_kv_cache_context( + is_paged_attn=True, + key_cache={0: key_cache}, + value_cache={0: value_cache}, + block_tables=torch.tensor([[0]], device=device, dtype=torch.int32), + context_lens=torch.tensor([seq_len], device=device, dtype=torch.int32), + slot_mapping=torch.arange(seq_len, device=device, dtype=torch.int64), + query_start_loc=torch.tensor([0, seq_len], device=device, dtype=torch.int32), + max_seqlen_q=seq_len, + ) + + module = torch.nn.Module() + module.layer_idx = 0 + actual, _ = flash_attention_forward( + module, + query, + key, + value, + attention_mask=None, + scaling=scaling, + sliding_window=sliding_window, + s_aux=sinks, + use_cute=True, + ) + + scores = torch.matmul(query.float(), key.float().transpose(-1, -2)) * scaling + positions = torch.arange(seq_len, device=device) + query_positions = positions[:, None] + key_positions = positions[None, :] + mask = (key_positions <= query_positions) & (key_positions > query_positions - sliding_window) + scores.masked_fill_(~mask, float("-inf")) + sink_scores = sinks.float().view(1, num_heads, 1, 1).expand(batch_size, -1, seq_len, -1) + attention_probs = torch.softmax(torch.cat((scores, sink_scores), dim=-1), dim=-1)[..., :-1] + expected = torch.matmul(attention_probs, value.float()) + expected = expected.transpose(1, 2).reshape_as(actual) + + torch.testing.assert_close(actual.float(), expected, rtol=2e-2, atol=2e-2) diff --git a/tests/test_model_config.py b/tests/test_model_config.py new file mode 100644 index 0000000..40da597 --- /dev/null +++ b/tests/test_model_config.py @@ -0,0 +1,34 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest + +from vexact.config import ModelConfig + + +def _gpt_oss_config(): + return SimpleNamespace(model_type="gpt_oss", max_position_embeddings=128) + + +def test_gpt_oss_requires_fa4_attention(): + with pytest.raises(ValueError, match="learnable attention sinks require"): + ModelConfig(model_path="unused", attn_impl="fa-invariant", hf_config=_gpt_oss_config()) + + +def test_gpt_oss_accepts_fa4_attention(): + config = ModelConfig(model_path="unused", attn_impl="fa-invariant-cute", hf_config=_gpt_oss_config()) + + assert config.attn_impl == "fa-invariant-cute" diff --git a/vexact/batch_invariant_ops/README.md b/vexact/batch_invariant_ops/README.md index 08fc0b3..f9d7619 100644 --- a/vexact/batch_invariant_ops/README.md +++ b/vexact/batch_invariant_ops/README.md @@ -61,7 +61,7 @@ batch_invariant_ops/ Two attention backends, both supporting paged KV cache: -- **Flash Attention** — Uses `flash_attn_with_kvcache` (FA3, SM90) or `flash_attn.cute` (FA4, SM90+/SM100). Both paths support causal sliding-window attention. Requires `is_paged_attn=True` in the KV cache context. Forces `num_splits=1` for determinism. +- **Flash Attention** — Uses `flash_attn_with_kvcache` (FA3, SM90) or `flash_attn.cute` (FA4, SM90+/SM100). Both paths support causal sliding-window attention; FA4 also supports per-head learnable attention sinks used by GPT-OSS. Requires `is_paged_attn=True` in the KV cache context. Forces `num_splits=1` for determinism. - **Flex Attention** — Uses `torch.nn.attention.flex_attention` with custom `mask_mod` for causal masking. Supports both packed (non-paged) and paged modes. Paged mode builds a physical-to-logical block mapping and caches the block mask across layers. Both expect the KV cache context to be set via `set_kv_cache_context()` before the forward pass. diff --git a/vexact/batch_invariant_ops/flash_attention.py b/vexact/batch_invariant_ops/flash_attention.py index 78b3884..6046ffa 100644 --- a/vexact/batch_invariant_ops/flash_attention.py +++ b/vexact/batch_invariant_ops/flash_attention.py @@ -35,6 +35,17 @@ def _get_fa_window_size(sliding_window: Optional[int]) -> tuple[int, int]: return (window_bound, window_bound) +def _get_fa_learnable_sink( + s_aux: Optional[torch.Tensor], query: torch.Tensor, use_cute: bool +) -> Optional[torch.Tensor]: + """Translate Transformers' attention sink argument to the supported FA backend.""" + if s_aux is None: + return None + if not use_cute: + raise NotImplementedError("Learnable attention sinks require the FA4 (fa-invariant-cute) backend") + return s_aux.to(dtype=query.dtype) + + @proton.scope("flash_attention_forward") def flash_attention_forward( module: nn.Module, @@ -45,6 +56,7 @@ def flash_attention_forward( scaling: float, dropout: float = 0.0, sliding_window: Optional[int] = None, + s_aux: Optional[torch.Tensor] = None, use_cute: bool = False, **kwargs, ): @@ -70,6 +82,7 @@ def flash_attention_forward( dropout (float): The dropout rate. Not currently used. sliding_window (Optional[int]): The number of tokens visible in causal local-attention layers, including the current token. ``None`` selects full causal attention. + s_aux (Optional[torch.Tensor]): Per-head learnable attention sink logits. Supported by FA4 only. Returns: Tuple[torch.Tensor, None]: A tuple containing the attention output and None. @@ -150,6 +163,7 @@ def flash_attention_forward( # max_seqlen_k = context_lens.max().item() from vexact.utils.device import DEVICE_MAJOR + learnable_sink = _get_fa_learnable_sink(s_aux, query, use_cute) if use_cute: assert DEVICE_MAJOR >= 9, f"FA4 (flash_attn.cute) requires SM90+, got SM{DEVICE_MAJOR}0" from flash_attn.cute import flash_attn_varlen_func @@ -168,6 +182,7 @@ def flash_attention_forward( softmax_scale=scaling, causal=True, window_size=window_size, + learnable_sink=learnable_sink, num_splits=1, ) else: diff --git a/vexact/config.py b/vexact/config.py index ff6c770..42e6b3e 100644 --- a/vexact/config.py +++ b/vexact/config.py @@ -143,6 +143,11 @@ def __post_init__(self): if self.hf_config is None: self._load_hf_config() + if getattr(self.hf_config, "model_type", None) == "gpt_oss" and self.attn_impl != "fa-invariant-cute": + raise ValueError( + f"GPT-OSS learnable attention sinks require attn_impl='fa-invariant-cute'; got {self.attn_impl!r}" + ) + # Initialize max_model_len self._initialize_max_model_len() From 09f2e72e7f0cada91bbd16d503cabe8da403718d Mon Sep 17 00:00:00 2001 From: "tianle.zhong" Date: Mon, 3 Aug 2026 23:49:05 +0000 Subject: [PATCH 2/3] [model] feat: Align GPT-OSS rollout MoE with VeOmni --- examples/moe/run_moonlight_gsm8k.sh | 1 + tests/test_model_config.py | 22 ++++++++++++++ tests/test_model_loader.py | 37 ++++++++++++++++++++++- vexact/config.py | 20 ++++++++++--- vexact/inferencer/model_loader.py | 46 ++++++++++++++++++++++++++--- vexact/worker/worker.py | 6 +++- 6 files changed, 122 insertions(+), 10 deletions(-) diff --git a/examples/moe/run_moonlight_gsm8k.sh b/examples/moe/run_moonlight_gsm8k.sh index d7912bf..2f5ca17 100755 --- a/examples/moe/run_moonlight_gsm8k.sh +++ b/examples/moe/run_moonlight_gsm8k.sh @@ -93,6 +93,7 @@ RAY_DEDUP_LOGS=1 PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \ actor_rollout_ref.rollout.max_num_batched_tokens=4096 \ ++actor_rollout_ref.rollout.engine_kwargs.vexact.max_cache_blocks=$vexact_max_cache_blocks \ ++actor_rollout_ref.rollout.engine_kwargs.vexact.attn_impl=$INFER_FA_IMPL \ + ++actor_rollout_ref.rollout.engine_kwargs.vexact.moe_implementation=$moe_implementation \ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \ actor_rollout_ref.rollout.disable_log_stats=False \ actor_rollout_ref.rollout.free_cache_engine=True \ diff --git a/tests/test_model_config.py b/tests/test_model_config.py index 40da597..6ceb8e0 100644 --- a/tests/test_model_config.py +++ b/tests/test_model_config.py @@ -32,3 +32,25 @@ def test_gpt_oss_accepts_fa4_attention(): config = ModelConfig(model_path="unused", attn_impl="fa-invariant-cute", hf_config=_gpt_oss_config()) assert config.attn_impl == "fa-invariant-cute" + + +def test_gpt_oss_rejects_non_quack_fused_moe(): + with pytest.raises(ValueError, match="requires moe_implementation='fused_quack'"): + ModelConfig( + model_path="unused", + attn_impl="fa-invariant-cute", + moe_implementation="fused_triton", + hf_config=_gpt_oss_config(), + ) + + +@pytest.mark.parametrize("moe_implementation", ["fused_quack", "eager"]) +def test_gpt_oss_accepts_supported_moe_implementations(moe_implementation): + config = ModelConfig( + model_path="unused", + attn_impl="fa-invariant-cute", + moe_implementation=moe_implementation, + hf_config=_gpt_oss_config(), + ) + + assert config.moe_implementation == moe_implementation diff --git a/tests/test_model_loader.py b/tests/test_model_loader.py index 8bc3dbe..de47533 100644 --- a/tests/test_model_loader.py +++ b/tests/test_model_loader.py @@ -13,17 +13,19 @@ # limitations under the License. import os +import sys import pytest import torch from torch import nn -from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig +from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel from vexact.config import PPInfo from vexact.inferencer.model_loader import ( ModelCreator, PPMissingLayer, TransformersForCausalLM, + _create_causal_model, load_weights_from_weight_iterator, ) @@ -191,3 +193,36 @@ def __init__(self): assert torch.equal(base_model.embed_tokens.weight, new_embed) assert torch.equal(base_model.proj.weight, new_proj) assert torch.equal(lm_head.weight, new_lm_head) + + +def test_gpt_oss_model_uses_veomni_modeling_and_moe_slot(monkeypatch): + class DummySlot: + def __init__(self): + self.impl = None + + def bind(self, impl): + self.impl = impl + + class DummyModel(PreTrainedModel): + config_class = PretrainedConfig + + def __init__(self, config): + super().__init__(config) + + slot = DummySlot() + modeling_module = sys.modules[__name__] + monkeypatch.setattr(modeling_module, "veomni_moe_experts_forward", slot, raising=False) + monkeypatch.setattr(DummyModel, "__module__", __name__) + monkeypatch.setattr( + "veomni.models.transformers.gpt_oss.register_gpt_oss_modeling", + lambda architecture: DummyModel, + ) + config = PretrainedConfig(architectures=["GptOssForCausalLM"]) + config.model_type = "gpt_oss" + config._attn_implementation = "fa-invariant-cute" + + model = _create_causal_model(config, "fused_quack") + + assert isinstance(model, DummyModel) + assert slot.impl == "quack" + assert model.config._attn_implementation == "fa-invariant-cute" diff --git a/vexact/config.py b/vexact/config.py index 42e6b3e..9178b26 100644 --- a/vexact/config.py +++ b/vexact/config.py @@ -100,6 +100,12 @@ class ModelConfig: default="fa-invariant", metadata={"help": "Attention implementation to use (fa-invariant, fa-invariant-cute, triton-invariant, flex)."}, ) + moe_implementation: str = field( + default="fused_quack", + metadata={ + "help": "MoE experts implementation shared with the VeOmni actor (fused_quack or eager for GPT-OSS)." + }, + ) enable_batch_invariant: bool = field( default=True, metadata={"help": "Enable batch invariant operations for deterministic inference."}, @@ -143,10 +149,16 @@ def __post_init__(self): if self.hf_config is None: self._load_hf_config() - if getattr(self.hf_config, "model_type", None) == "gpt_oss" and self.attn_impl != "fa-invariant-cute": - raise ValueError( - f"GPT-OSS learnable attention sinks require attn_impl='fa-invariant-cute'; got {self.attn_impl!r}" - ) + if getattr(self.hf_config, "model_type", None) == "gpt_oss": + if self.attn_impl != "fa-invariant-cute": + raise ValueError( + f"GPT-OSS learnable attention sinks require attn_impl='fa-invariant-cute'; got {self.attn_impl!r}" + ) + if self.moe_implementation not in ("fused_quack", "eager"): + raise ValueError( + "GPT-OSS requires moe_implementation='fused_quack' (SM90+) or 'eager'; " + f"got {self.moe_implementation!r}" + ) # Initialize max_model_len self._initialize_max_model_len() diff --git a/vexact/inferencer/model_loader.py b/vexact/inferencer/model_loader.py index 5ac1921..ff434fe 100644 --- a/vexact/inferencer/model_loader.py +++ b/vexact/inferencer/model_loader.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import sys from contextlib import contextmanager from dataclasses import dataclass from typing import Generator, Iterable, Optional @@ -36,6 +37,38 @@ logger = logging.getLogger(__name__) +def _create_causal_model(config: PretrainedConfig, moe_implementation: str) -> PreTrainedModel: + """Construct a causal model, using VeOmni's GPT-OSS modeling and MoE dispatch when required.""" + if getattr(config, "model_type", None) != "gpt_oss": + return AutoModelForCausalLM.from_config(config) + + from veomni.models.transformers.gpt_oss import register_gpt_oss_modeling + + architectures = getattr(config, "architectures", None) or ["GptOssForCausalLM"] + architecture = architectures[0] if isinstance(architectures, (list, tuple)) else architectures + model_cls = register_gpt_oss_modeling(architecture) + modeling_module = sys.modules[model_cls.__module__] + + # VeOmni config values use a ``fused_`` prefix while OpSlot registry names do not. + # Binding the same choice as the actor is what makes strict rollout/training + # probability alignment meaningful for GPT-OSS's interleaved gate/up experts. + moe_impl = "eager" if moe_implementation == "eager" else moe_implementation.removeprefix("fused_") + modeling_module.veomni_moe_experts_forward.bind(moe_impl) + logger.info("GPT-OSS rollout uses VeOmni modeling with MoE implementation %s", moe_implementation) + + # Transformers validates attention backend names in PreTrainedModel.__init__, + # before Vexact's custom dispatcher is consulted. GPT-OSS uses runtime + # attention dispatch, so initialize its unchanged module structure with an + # accepted backend name and restore Vexact's backend immediately afterward. + attn_implementation = config._attn_implementation + config._attn_implementation = "eager" + try: + model = model_cls(config) + finally: + config._attn_implementation = attn_implementation + return model + + @dataclass class StateDictIterator: filepaths: list[str] @@ -277,15 +310,20 @@ def forward( class ModelCreator: - def __init__(self, config: PretrainedConfig, model_path: str, device: torch.device, pp_info: PPInfo): + def __init__( + self, + config: PretrainedConfig, + model_path: str, + device: torch.device, + pp_info: PPInfo, + moe_implementation: str = "fused_quack", + ): self._pp_info = pp_info self._model_path = model_path self._config = config self._device = device with init_on_device_without_buffers("meta"): - self._causal_model: PreTrainedModel = AutoModelForCausalLM.from_config( - self._config, - ) + self._causal_model = _create_causal_model(self._config, moe_implementation) self.model = self._causal_model.model def create_model(self): diff --git a/vexact/worker/worker.py b/vexact/worker/worker.py index c4c303d..f71e48a 100644 --- a/vexact/worker/worker.py +++ b/vexact/worker/worker.py @@ -102,7 +102,11 @@ def __init__(self, config: VeXactConfig, rank: int): # Load model self.model = ModelCreator( - self.config.model.hf_config, config.model.model_path, self.device, self.pp_info + self.config.model.hf_config, + config.model.model_path, + self.device, + self.pp_info, + moe_implementation=config.model.moe_implementation, ).create_model() # Create PP messager if needed From 0cfb5bbd56e303823870cdf56815ff28b637eefb Mon Sep 17 00:00:00 2001 From: "tianle.zhong" Date: Mon, 3 Aug 2026 23:56:27 +0000 Subject: [PATCH 3/3] [model] refactor: Isolate GPT-OSS MoE patch --- tests/test_gpt_oss_model.py | 92 ++++++++++++++++++++++ tests/test_model_config.py | 12 +-- tests/test_model_loader.py | 37 +-------- vexact/config.py | 9 +-- vexact/inferencer/model_loader.py | 46 +---------- vexact/models/gpt_oss/__init__.py | 18 +++++ vexact/models/gpt_oss/modeling_gpt_oss.py | 93 +++++++++++++++++++++++ vexact/models/register.py | 5 ++ vexact/worker/worker.py | 6 +- 9 files changed, 223 insertions(+), 95 deletions(-) create mode 100644 tests/test_gpt_oss_model.py create mode 100644 vexact/models/gpt_oss/__init__.py create mode 100644 vexact/models/gpt_oss/modeling_gpt_oss.py diff --git a/tests/test_gpt_oss_model.py b/tests/test_gpt_oss_model.py new file mode 100644 index 0000000..674dc71 --- /dev/null +++ b/tests/test_gpt_oss_model.py @@ -0,0 +1,92 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +import torch +from transformers import AutoModelForCausalLM, GptOssConfig +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts, GptOssForCausalLM + +from vexact.inferencer.model_loader import init_on_device_without_buffers +from vexact.models.gpt_oss import modeling_gpt_oss + + +def test_gpt_oss_experts_use_fused_quack(monkeypatch): + captured = {} + + def fake_fused_quack(**kwargs): + captured.update(kwargs) + return torch.tensor([1.0]) + + monkeypatch.setattr(modeling_gpt_oss, "_run_fused_quack", fake_fused_quack) + experts = SimpleNamespace( + num_experts=4, + gate_up_proj=torch.tensor([1.0]), + gate_up_proj_bias=torch.tensor([2.0]), + down_proj=torch.tensor([3.0]), + down_proj_bias=torch.tensor([4.0]), + alpha=1.702, + limit=7.0, + ) + hidden_states = torch.tensor([[5.0]]) + router_indices = torch.tensor([[1]]) + routing_weights = torch.tensor([[0.5]]) + + output = modeling_gpt_oss.gpt_oss_experts_forward(experts, hidden_states, router_indices, routing_weights) + + assert torch.equal(output, torch.tensor([1.0])) + assert captured == { + "num_experts": 4, + "routing_weights": routing_weights, + "selected_experts": router_indices, + "hidden_states": hidden_states, + "gate_up_proj": experts.gate_up_proj, + "gate_up_proj_bias": experts.gate_up_proj_bias, + "down_proj": experts.down_proj, + "down_proj_bias": experts.down_proj_bias, + "alpha": 1.702, + "limit": 7.0, + } + + +def test_gpt_oss_experts_require_routing_inputs(): + with pytest.raises(ValueError, match="requires router_indices and routing_weights"): + modeling_gpt_oss.gpt_oss_experts_forward(SimpleNamespace(), torch.tensor([[1.0]])) + + +def test_gpt_oss_patch_replaces_experts_forward(): + modeling_gpt_oss.apply_gpt_oss_patches() + + assert GptOssExperts.forward is modeling_gpt_oss.gpt_oss_experts_forward + + +def test_gpt_oss_uses_common_transformers_model_path(): + config = GptOssConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + num_local_experts=4, + num_experts_per_tok=2, + head_dim=16, + ) + with init_on_device_without_buffers("meta"): + model = AutoModelForCausalLM.from_config(config) + + assert isinstance(model, GptOssForCausalLM) + assert type(model).__module__.startswith("transformers.models.gpt_oss") + assert model.model.layers[0].mlp.experts.forward.__func__ is modeling_gpt_oss.gpt_oss_experts_forward diff --git a/tests/test_model_config.py b/tests/test_model_config.py index 6ceb8e0..a61d8af 100644 --- a/tests/test_model_config.py +++ b/tests/test_model_config.py @@ -34,23 +34,23 @@ def test_gpt_oss_accepts_fa4_attention(): assert config.attn_impl == "fa-invariant-cute" -def test_gpt_oss_rejects_non_quack_fused_moe(): +@pytest.mark.parametrize("moe_implementation", ["fused_triton", "eager"]) +def test_gpt_oss_rejects_non_quack_moe(moe_implementation): with pytest.raises(ValueError, match="requires moe_implementation='fused_quack'"): ModelConfig( model_path="unused", attn_impl="fa-invariant-cute", - moe_implementation="fused_triton", + moe_implementation=moe_implementation, hf_config=_gpt_oss_config(), ) -@pytest.mark.parametrize("moe_implementation", ["fused_quack", "eager"]) -def test_gpt_oss_accepts_supported_moe_implementations(moe_implementation): +def test_gpt_oss_accepts_fused_quack_moe(): config = ModelConfig( model_path="unused", attn_impl="fa-invariant-cute", - moe_implementation=moe_implementation, + moe_implementation="fused_quack", hf_config=_gpt_oss_config(), ) - assert config.moe_implementation == moe_implementation + assert config.moe_implementation == "fused_quack" diff --git a/tests/test_model_loader.py b/tests/test_model_loader.py index de47533..8bc3dbe 100644 --- a/tests/test_model_loader.py +++ b/tests/test_model_loader.py @@ -13,19 +13,17 @@ # limitations under the License. import os -import sys import pytest import torch from torch import nn -from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel +from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig from vexact.config import PPInfo from vexact.inferencer.model_loader import ( ModelCreator, PPMissingLayer, TransformersForCausalLM, - _create_causal_model, load_weights_from_weight_iterator, ) @@ -193,36 +191,3 @@ def __init__(self): assert torch.equal(base_model.embed_tokens.weight, new_embed) assert torch.equal(base_model.proj.weight, new_proj) assert torch.equal(lm_head.weight, new_lm_head) - - -def test_gpt_oss_model_uses_veomni_modeling_and_moe_slot(monkeypatch): - class DummySlot: - def __init__(self): - self.impl = None - - def bind(self, impl): - self.impl = impl - - class DummyModel(PreTrainedModel): - config_class = PretrainedConfig - - def __init__(self, config): - super().__init__(config) - - slot = DummySlot() - modeling_module = sys.modules[__name__] - monkeypatch.setattr(modeling_module, "veomni_moe_experts_forward", slot, raising=False) - monkeypatch.setattr(DummyModel, "__module__", __name__) - monkeypatch.setattr( - "veomni.models.transformers.gpt_oss.register_gpt_oss_modeling", - lambda architecture: DummyModel, - ) - config = PretrainedConfig(architectures=["GptOssForCausalLM"]) - config.model_type = "gpt_oss" - config._attn_implementation = "fa-invariant-cute" - - model = _create_causal_model(config, "fused_quack") - - assert isinstance(model, DummyModel) - assert slot.impl == "quack" - assert model.config._attn_implementation == "fa-invariant-cute" diff --git a/vexact/config.py b/vexact/config.py index 9178b26..4d20589 100644 --- a/vexact/config.py +++ b/vexact/config.py @@ -102,9 +102,7 @@ class ModelConfig: ) moe_implementation: str = field( default="fused_quack", - metadata={ - "help": "MoE experts implementation shared with the VeOmni actor (fused_quack or eager for GPT-OSS)." - }, + metadata={"help": "MoE experts implementation. GPT-OSS supports fused_quack only."}, ) enable_batch_invariant: bool = field( default=True, @@ -154,10 +152,9 @@ def __post_init__(self): raise ValueError( f"GPT-OSS learnable attention sinks require attn_impl='fa-invariant-cute'; got {self.attn_impl!r}" ) - if self.moe_implementation not in ("fused_quack", "eager"): + if self.moe_implementation != "fused_quack": raise ValueError( - "GPT-OSS requires moe_implementation='fused_quack' (SM90+) or 'eager'; " - f"got {self.moe_implementation!r}" + f"GPT-OSS requires moe_implementation='fused_quack' (SM90+); got {self.moe_implementation!r}" ) # Initialize max_model_len diff --git a/vexact/inferencer/model_loader.py b/vexact/inferencer/model_loader.py index ff434fe..5ac1921 100644 --- a/vexact/inferencer/model_loader.py +++ b/vexact/inferencer/model_loader.py @@ -13,7 +13,6 @@ # limitations under the License. import logging -import sys from contextlib import contextmanager from dataclasses import dataclass from typing import Generator, Iterable, Optional @@ -37,38 +36,6 @@ logger = logging.getLogger(__name__) -def _create_causal_model(config: PretrainedConfig, moe_implementation: str) -> PreTrainedModel: - """Construct a causal model, using VeOmni's GPT-OSS modeling and MoE dispatch when required.""" - if getattr(config, "model_type", None) != "gpt_oss": - return AutoModelForCausalLM.from_config(config) - - from veomni.models.transformers.gpt_oss import register_gpt_oss_modeling - - architectures = getattr(config, "architectures", None) or ["GptOssForCausalLM"] - architecture = architectures[0] if isinstance(architectures, (list, tuple)) else architectures - model_cls = register_gpt_oss_modeling(architecture) - modeling_module = sys.modules[model_cls.__module__] - - # VeOmni config values use a ``fused_`` prefix while OpSlot registry names do not. - # Binding the same choice as the actor is what makes strict rollout/training - # probability alignment meaningful for GPT-OSS's interleaved gate/up experts. - moe_impl = "eager" if moe_implementation == "eager" else moe_implementation.removeprefix("fused_") - modeling_module.veomni_moe_experts_forward.bind(moe_impl) - logger.info("GPT-OSS rollout uses VeOmni modeling with MoE implementation %s", moe_implementation) - - # Transformers validates attention backend names in PreTrainedModel.__init__, - # before Vexact's custom dispatcher is consulted. GPT-OSS uses runtime - # attention dispatch, so initialize its unchanged module structure with an - # accepted backend name and restore Vexact's backend immediately afterward. - attn_implementation = config._attn_implementation - config._attn_implementation = "eager" - try: - model = model_cls(config) - finally: - config._attn_implementation = attn_implementation - return model - - @dataclass class StateDictIterator: filepaths: list[str] @@ -310,20 +277,15 @@ def forward( class ModelCreator: - def __init__( - self, - config: PretrainedConfig, - model_path: str, - device: torch.device, - pp_info: PPInfo, - moe_implementation: str = "fused_quack", - ): + def __init__(self, config: PretrainedConfig, model_path: str, device: torch.device, pp_info: PPInfo): self._pp_info = pp_info self._model_path = model_path self._config = config self._device = device with init_on_device_without_buffers("meta"): - self._causal_model = _create_causal_model(self._config, moe_implementation) + self._causal_model: PreTrainedModel = AutoModelForCausalLM.from_config( + self._config, + ) self.model = self._causal_model.model def create_model(self): diff --git a/vexact/models/gpt_oss/__init__.py b/vexact/models/gpt_oss/__init__.py new file mode 100644 index 0000000..1315341 --- /dev/null +++ b/vexact/models/gpt_oss/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .modeling_gpt_oss import apply_gpt_oss_patches + + +__all__ = ["apply_gpt_oss_patches"] diff --git a/vexact/models/gpt_oss/modeling_gpt_oss.py b/vexact/models/gpt_oss/modeling_gpt_oss.py new file mode 100644 index 0000000..ae16d04 --- /dev/null +++ b/vexact/models/gpt_oss/modeling_gpt_oss.py @@ -0,0 +1,93 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Patch HuggingFace GPT-OSS experts to use VeOmni's fused quack kernel.""" + +import logging + +import torch +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + + +logger = logging.getLogger(__name__) +_PATCHED = False + + +def _run_fused_quack( + *, + num_experts: int, + routing_weights: torch.Tensor, + selected_experts: torch.Tensor, + hidden_states: torch.Tensor, + gate_up_proj: torch.Tensor, + gate_up_proj_bias: torch.Tensor, + down_proj: torch.Tensor, + down_proj_bias: torch.Tensor, + alpha: float, + limit: float, +) -> torch.Tensor: + # Import lazily: register_models() also runs in CPU-only verl workers, while + # the quack kernel itself is only available on SM90+ rollout workers. + from veomni.ops.kernels.moe.quack_gemm_interleave_gate_up import ( + quack_gemm_gpt_oss_fused_moe_forward, + ) + + return quack_gemm_gpt_oss_fused_moe_forward( + num_experts=num_experts, + routing_weights=routing_weights, + selected_experts=selected_experts, + hidden_states=hidden_states, + gate_up_proj=gate_up_proj, + gate_up_proj_bias=gate_up_proj_bias, + down_proj=down_proj, + down_proj_bias=down_proj_bias, + alpha=alpha, + limit=limit, + ) + + +def gpt_oss_experts_forward( + self: GptOssExperts, + hidden_states: torch.Tensor, + router_indices: torch.Tensor | None = None, + routing_weights: torch.Tensor | None = None, +) -> torch.Tensor: + if router_indices is None or routing_weights is None: + raise ValueError("GPT-OSS fused_quack requires router_indices and routing_weights") + + return _run_fused_quack( + num_experts=self.num_experts, + routing_weights=routing_weights, + selected_experts=router_indices, + hidden_states=hidden_states, + gate_up_proj=self.gate_up_proj, + gate_up_proj_bias=self.gate_up_proj_bias, + down_proj=self.down_proj, + down_proj_bias=self.down_proj_bias, + alpha=self.alpha, + limit=self.limit, + ) + + +def apply_gpt_oss_patches() -> None: + global _PATCHED + if _PATCHED: + return + + GptOssExperts.forward = gpt_oss_experts_forward + logger.info("Applied GPT-OSS fused_quack monkey patch.") + _PATCHED = True + + +__all__ = ["apply_gpt_oss_patches", "gpt_oss_experts_forward"] diff --git a/vexact/models/register.py b/vexact/models/register.py index ff4a818..c745f26 100644 --- a/vexact/models/register.py +++ b/vexact/models/register.py @@ -66,6 +66,11 @@ def register_models() -> None: apply_deepseek_v3_patches() print("[VEXACT] register_models(): DeepSeek-V3 monkey patch active") + from .gpt_oss.modeling_gpt_oss import apply_gpt_oss_patches + + apply_gpt_oss_patches() + print("[VEXACT] register_models(): GPT-OSS fused_quack monkey patch active") + use_liger_patch = os.getenv("VEXACT_LIGER_PATCH", "0") == "1" if use_liger_patch: # This path mimics VeOmni Qwen3 GPU patch to use Liger kernels diff --git a/vexact/worker/worker.py b/vexact/worker/worker.py index f71e48a..c4c303d 100644 --- a/vexact/worker/worker.py +++ b/vexact/worker/worker.py @@ -102,11 +102,7 @@ def __init__(self, config: VeXactConfig, rank: int): # Load model self.model = ModelCreator( - self.config.model.hf_config, - config.model.model_path, - self.device, - self.pp_info, - moe_implementation=config.model.moe_implementation, + self.config.model.hf_config, config.model.model_path, self.device, self.pp_info ).create_model() # Create PP messager if needed