From 3bc19633cb921cce1453708f1614b36096473bf6 Mon Sep 17 00:00:00 2001 From: DAN_111111 Date: Sun, 2 Aug 2026 11:56:44 +0000 Subject: [PATCH] Update preset code to support new versions of transformers --- roundpipe/models/gpt_oss.py | 44 ++++- roundpipe/models/qwen3.py | 79 +++++++-- tests/models/test_preset_compatibility.py | 189 ++++++++++++++++++++++ 3 files changed, 290 insertions(+), 22 deletions(-) create mode 100644 tests/models/test_preset_compatibility.py diff --git a/roundpipe/models/gpt_oss.py b/roundpipe/models/gpt_oss.py index f89334a..35bb312 100644 --- a/roundpipe/models/gpt_oss.py +++ b/roundpipe/models/gpt_oss.py @@ -20,9 +20,30 @@ from .function import CompileForCausalLMLoss, ChunkedCompileLinearForCausalLMLoss +def _resolve_attention_type(layer: GptOssDecoderLayer) -> str: + attention_type = getattr(layer.self_attn, "layer_type", None) + if attention_type is None: + attention_type = getattr(layer.self_attn, "attention_type", None) + if attention_type is None: + attention_type = getattr(layer, "attention_type", None) + if attention_type is None: + raise AttributeError("GPT-OSS decoder layer does not expose its attention type") + return cast(str, attention_type) + + class GptOssOptExperts(nn.Module): def __init__(self, mod: GptOssExperts) -> None: super().__init__() + expert_weights = (mod.gate_up_proj, mod.down_proj) + if any( + weight.ndim != 3 or not torch.is_floating_point(weight) + for weight in expert_weights + ): + raise TypeError( + "RoundPipe's GPT-OSS preset requires unpacked floating-point " + "expert weights; load a BF16/FP16 checkpoint instead of an " + "MXFP4/packed checkpoint." + ) self.num_experts = mod.num_experts self.hidden_size = mod.hidden_size self.alpha = mod.alpha @@ -39,13 +60,20 @@ def forward( router_indices: torch.Tensor, routing_weights: torch.Tensor, ) -> torch.Tensor: - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) + input_shape = hidden_states.shape + hidden_dim = input_shape[-1] + hidden_states = hidden_states.reshape(-1, hidden_dim) top_k = router_indices.shape[-1] - selected_experts = router_indices.view(-1) - routing_weights = torch.gather(routing_weights, 1, router_indices) - routing_weights = routing_weights.view(-1) + selected_experts = router_indices.reshape(-1) + if routing_weights.shape != router_indices.shape: + if routing_weights.shape[-1] != self.num_experts: + raise ValueError( + "GPT-OSS routing weights must contain either top-k weights " + "or one weight per expert" + ) + routing_weights = torch.gather(routing_weights, 1, router_indices) + routing_weights = routing_weights.reshape(-1) _, sort_idx = torch.sort(selected_experts) permute_weight = routing_weights[sort_idx] @@ -66,7 +94,7 @@ def forward( save_for_recompute(token_per_expert_cpu) final_hidden_states = torch.zeros( - (batch_size * sequence_length, hidden_dim), + (hidden_states.shape[0], hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) @@ -97,7 +125,7 @@ def forward( final_hidden_states.index_add_(0, expert_tokens, expert_output) start_idx += num_tokens - return final_hidden_states.view(batch_size, sequence_length, hidden_dim) + return final_hidden_states.reshape(input_shape) class GptOssForCausalLMPrefix(nn.Module): @@ -203,7 +231,7 @@ def __init__(self, layer: GptOssDecoderLayer) -> None: self.mlp = layer.mlp self.input_layernorm = layer.input_layernorm self.post_attention_layernorm = layer.post_attention_layernorm - self.attention_type = layer.self_attn.attention_type + self.attention_type = _resolve_attention_type(layer) def forward(self, input): ( diff --git a/roundpipe/models/qwen3.py b/roundpipe/models/qwen3.py index 92dc654..44dc1f3 100644 --- a/roundpipe/models/qwen3.py +++ b/roundpipe/models/qwen3.py @@ -1,4 +1,5 @@ from typing_extensions import * +import inspect import warnings import torch @@ -15,6 +16,45 @@ from .function import CompileForCausalLMLoss, ChunkedCompileLinearForCausalLMLoss +def _resolve_attention_type(layer: nn.Module) -> str: + self_attn = layer.self_attn + attention_type = getattr(self_attn, "layer_type", None) + if attention_type is None: + attention_type = getattr(self_attn, "attention_type", None) + if attention_type is None: + attention_type = getattr(layer, "attention_type", None) + if attention_type is None: + raise AttributeError("Qwen3 decoder layer does not expose its attention type") + return cast(str, attention_type) + + +def _create_mask_compat( + mask_function: Callable[..., torch.Tensor], + *, + config: Any, + inputs_embeds: torch.Tensor, + attention_mask: Optional[torch.Tensor], + past_key_values: Optional[Any], + position_ids: torch.Tensor, + cache_position: Optional[torch.Tensor], +) -> torch.Tensor: + parameters = inspect.signature(mask_function).parameters + mask_kwargs: Dict[str, Any] = { + "config": config, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + } + embeds_parameter = ( + "inputs_embeds" if "inputs_embeds" in parameters else "input_embeds" + ) + mask_kwargs[embeds_parameter] = inputs_embeds + if "position_ids" in parameters: + mask_kwargs["position_ids"] = position_ids + if "cache_position" in parameters: + mask_kwargs["cache_position"] = cache_position + return mask_function(**mask_kwargs) + + class Qwen3ForCausalLMPrefix(nn.Module): def __init__(self, model: Qwen3ForCausalLM) -> None: super().__init__() @@ -70,29 +110,39 @@ def forward( ) past_key_values = None - if position_ids is None: - position_ids = torch.arange( + if cache_position is None: + cache_position = torch.arange( inputs_embeds.shape[1], device=inputs_embeds.device - ).unsqueeze(0) + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) # It may already have been prepared by e.g. `generate` if not isinstance(causal_mask_mapping := attention_mask, dict): - # Prepare mask arguments - mask_kwargs = { - "config": self.config, - "inputs_embeds": inputs_embeds, - "attention_mask": attention_mask, - "past_key_values": past_key_values, - "position_ids": position_ids, - } # Create the masks causal_mask_mapping = { - "full_attention": create_causal_mask(**mask_kwargs), + "full_attention": _create_mask_compat( + create_causal_mask, + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + cache_position=cache_position, + ), } # The sliding window alternating layers are not always activated depending on the config if self.has_sliding_layers: causal_mask_mapping["sliding_attention"] = ( - create_sliding_window_causal_mask(**mask_kwargs) + _create_mask_compat( + create_sliding_window_causal_mask, + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + cache_position=cache_position, + ) ) hidden_states = inputs_embeds @@ -114,6 +164,7 @@ class Qwen3ForCausalLMWrappedLayer(nn.Module): def __init__(self, layer: nn.Module) -> None: super().__init__() self.layer = layer + self.attention_type = _resolve_attention_type(layer) def forward(self, input): ( @@ -127,7 +178,7 @@ def forward(self, input): ) = input hidden_states = self.layer( hidden_states, - attention_mask=causal_mask_mapping[self.layer.self_attn.layer_type], + attention_mask=causal_mask_mapping[self.attention_type], position_ids=position_ids, past_key_values=None, use_cache=False, diff --git a/tests/models/test_preset_compatibility.py b/tests/models/test_preset_compatibility.py new file mode 100644 index 0000000..4231966 --- /dev/null +++ b/tests/models/test_preset_compatibility.py @@ -0,0 +1,189 @@ +from types import SimpleNamespace + +import pytest +import torch +from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + +from roundpipe.models import gpt_oss, qwen3 + + +def test_qwen3_resolves_current_and_legacy_attention_type_locations(): + current_layer = SimpleNamespace( + self_attn=SimpleNamespace(layer_type="sliding_attention") + ) + legacy_layer = SimpleNamespace( + self_attn=SimpleNamespace(), attention_type="full_attention" + ) + + assert qwen3._resolve_attention_type(current_layer) == "sliding_attention" + assert qwen3._resolve_attention_type(legacy_layer) == "full_attention" + + +def test_qwen3_calls_current_mask_contract(): + inputs_embeds = torch.randn(1, 3, 4) + position_ids = torch.arange(3).unsqueeze(0) + + def current_mask( + *, + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + ): + return ( + config, + inputs_embeds, + attention_mask, + past_key_values, + position_ids, + ) + + result = qwen3._create_mask_compat( + current_mask, + config="config", + inputs_embeds=inputs_embeds, + attention_mask="mask", + past_key_values=None, + position_ids=position_ids, + cache_position=None, + ) + + assert result == ("config", inputs_embeds, "mask", None, position_ids) + + +def test_qwen3_calls_legacy_mask_contract(): + inputs_embeds = torch.randn(1, 3, 4) + position_ids = torch.arange(3).unsqueeze(0) + cache_position = position_ids.squeeze(0) + + def legacy_mask( + *, + config, + input_embeds, + attention_mask, + cache_position, + past_key_values, + position_ids, + ): + return ( + config, + input_embeds, + attention_mask, + cache_position, + past_key_values, + position_ids, + ) + + result = qwen3._create_mask_compat( + legacy_mask, + config="config", + inputs_embeds=inputs_embeds, + attention_mask="mask", + past_key_values=None, + position_ids=position_ids, + cache_position=cache_position, + ) + + assert result == ( + "config", + inputs_embeds, + "mask", + cache_position, + None, + position_ids, + ) + + +def _make_gpt_oss_experts() -> GptOssExperts: + config = GptOssConfig( + hidden_size=8, + intermediate_size=4, + num_local_experts=3, + num_experts_per_tok=2, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + vocab_size=32, + ) + torch.manual_seed(0) + experts = GptOssExperts(config) + for parameter in experts.parameters(): + torch.nn.init.uniform_(parameter, -0.1, 0.1) + return experts + + +def test_gpt_oss_optimized_experts_support_current_2d_topk_contract(): + experts = _make_gpt_oss_experts() + optimized = gpt_oss.GptOssOptExperts(experts) + router_indices = torch.tensor([[0, 2], [1, 0], [2, 1]]) + routing_weights = torch.tensor([[0.7, 0.3], [0.4, 0.6], [0.8, 0.2]]) + reference_input = torch.randn(3, 8, requires_grad=True) + optimized_input = reference_input.detach().clone().requires_grad_(True) + + expected = experts(reference_input, router_indices, routing_weights) + actual = optimized(optimized_input, router_indices, routing_weights) + + torch.testing.assert_close(actual, expected) + expected.sum().backward() + actual.sum().backward() + torch.testing.assert_close(optimized_input.grad, reference_input.grad) + + +def test_gpt_oss_optimized_experts_support_legacy_3d_full_router_contract(): + experts = _make_gpt_oss_experts() + optimized = gpt_oss.GptOssOptExperts(experts) + hidden_states = torch.randn(2, 2, 8) + router_indices = torch.tensor([[0, 2], [1, 0], [2, 1], [1, 2]]) + full_routing_weights = torch.tensor( + [ + [0.7, 0.0, 0.3], + [0.6, 0.4, 0.0], + [0.0, 0.2, 0.8], + [0.0, 0.4, 0.6], + ] + ) + topk_routing_weights = torch.gather( + full_routing_weights, 1, router_indices + ) + + expected = experts( + hidden_states.reshape(-1, 8), router_indices, topk_routing_weights + ).reshape_as(hidden_states) + actual = optimized(hidden_states, router_indices, full_routing_weights) + + torch.testing.assert_close(actual, expected) + + +def test_gpt_oss_resolves_current_and_legacy_attention_type_locations(): + current_layer = SimpleNamespace( + self_attn=SimpleNamespace(layer_type="sliding_attention") + ) + intermediary_layer = SimpleNamespace( + self_attn=SimpleNamespace(attention_type="full_attention") + ) + legacy_layer = SimpleNamespace( + self_attn=SimpleNamespace(), attention_type="sliding_attention" + ) + + assert gpt_oss._resolve_attention_type(current_layer) == "sliding_attention" + assert gpt_oss._resolve_attention_type(intermediary_layer) == "full_attention" + assert gpt_oss._resolve_attention_type(legacy_layer) == "sliding_attention" + + +def test_gpt_oss_optimized_experts_reject_packed_weights_at_wrap_time(): + packed_experts = SimpleNamespace( + num_experts=3, + hidden_size=8, + alpha=1.702, + limit=7.0, + gate_up_proj=torch.empty(3, 8, 2, dtype=torch.uint8), + gate_up_proj_bias=torch.empty(3, 8), + down_proj=torch.empty(3, 4, 2, dtype=torch.uint8), + down_proj_bias=torch.empty(3, 8), + ) + + with pytest.raises(TypeError, match="floating-point expert weights"): + gpt_oss.GptOssOptExperts(packed_experts)