diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 45606dd532..fcfeae1bcc 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -597,6 +597,10 @@ jobs: "num_gpus": 0, "test_file": "test_cispo_loss.py" }, + { + "num_gpus": 0, + "test_file": "test_qwen3_5_vl_gdn_packed.py" + }, { "num_gpus": 0, "test_file": "test_ppo_logprob_entropy.py" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index f58b064830..f92ffa9809 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -74,6 +74,7 @@ {'test_file': 'test_logprob_response_spans.py', 'num_gpus': 0}, {'test_file': 'test_value_temperature.py', 'num_gpus': 0}, {'test_file': 'test_cispo_loss.py', 'num_gpus': 0}, + {'test_file': 'test_qwen3_5_vl_gdn_packed.py', 'num_gpus': 0}, {'test_file': 'test_ppo_logprob_entropy.py', 'num_gpus': 0}, {'test_file': 'test_rm_f1.py', 'num_gpus': 0}, {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0}, diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md index d5170617de..38206c9902 100644 --- a/examples/geo3k_vlm/README.md +++ b/examples/geo3k_vlm/README.md @@ -12,12 +12,7 @@ Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prev pip install nvidia-cudnn-cu12==9.16.0.29 ``` -**Important:** We use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to support multimodal models. However, not all Megatron arguments are passed through to Megatron Bridge — you may need to set some manually [here](https://github.com/THUDM/slime/blob/de84e10d468dcb726e1199fd6bd16aa9538aed09/slime/backends/megatron_utils/model_provider.py#L89) (currently only parallelization-related arguments are passed). For example, for Qwen3-VL-30B-A3B you may need to add: -```python -provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff -provider.freeze_language_model = False -provider.freeze_vision_model = False -``` +**Important:** We use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to support multimodal models. Besides the parallelization arguments, slime forwards `--moe-aux-loss-coeff`, `--freeze-language-model`, `--freeze-vision-model` and `--freeze-vision-projection` to the provider when the provider defines them. Arguments outside that list still need to be set by hand in `slime/backends/megatron_utils/model_provider.py`.

Reward Plot @@ -95,7 +90,9 @@ SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm. #### Qwen3.5 Series We provide an [example](./run_geo3k_qwen35.sh) for Qwen3.5-35B-A3B. To support other Qwen3.5 models, add a model config file in `scripts/models/` and update the model name and config path in the script accordingly. -For GDN training, use `--micro-batch-size 1` and remove `--use-dynamic-batch-size`. +For GDN training, use `--micro-batch-size 1` and remove `--use-dynamic-batch-size`. Both are enforced: GDN runs one recurrence over the whole microbatch, so several packed sequences would be fused into one stream. Context parallel is rejected for the same reason. + +Keep `--attention-backend flash` as the example sets it. Letting Megatron pick the backend selects Transformer Engine's cuDNN fused attention, whose backward produces non-finite gradients for this model under bf16 with packed sequences and tensor parallel above 1. ## Notes diff --git a/examples/geo3k_vlm/run_geo3k_qwen35.sh b/examples/geo3k_vlm/run_geo3k_qwen35.sh index a53440cc50..6459377408 100644 --- a/examples/geo3k_vlm/run_geo3k_qwen35.sh +++ b/examples/geo3k_vlm/run_geo3k_qwen35.sh @@ -4,12 +4,10 @@ pip install -U transformers -# IMPORTANT: This branch is specially modified for slime's current Megatron -# version and Qwen3.5 from the main Megatron Bridge. Other models are not verified! -# To restore the original Megatron Bridge, run: -# pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation -# TODO: Remove this once Megatron & Megatron Bridge are upgraded upstream. -pip install git+https://github.com/coding-famer/Megatron-Bridge-slime.git@qwen35 --no-build-isolation +# Qwen3.5-VL is provided by Megatron Bridge itself since v0.4.0, which the slime +# image already ships, so no forked Megatron Bridge is needed here. +# slime_plugins/megatron_bridge/qwen3_5_vl.py adapts its GDN layers to the packed +# sequences slime feeds the model. # Configuration TRAIN_BACKEND="megatron" diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index 090b692f18..e1cb0e01d2 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/slime/backends/megatron_utils/model_provider.py @@ -20,6 +20,18 @@ from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config from slime.utils.misc import load_function +# Args that Megatron Bridge providers accept but that are not part of the +# parallelism config forwarded above. Previously users had to edit this file by +# hand to set them (see examples/geo3k_vlm/README.md). Only names present on +# both the CLI args and the provider are forwarded, so this stays a no-op for +# providers that do not define them. +_BRIDGE_PROVIDER_PASSTHROUGH_ARGS = ( + "moe_aux_loss_coeff", + "freeze_language_model", + "freeze_vision_model", + "freeze_vision_projection", +) + # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82 class LinearForLastLayer(torch.nn.Linear): @@ -105,6 +117,10 @@ def wrapped_model_provider( provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers if getattr(args, "decoder_last_pipeline_num_layers", None) is not None: provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers + for name in _BRIDGE_PROVIDER_PASSTHROUGH_ARGS: + value = getattr(args, name, None) + if value is not None and hasattr(provider, name): + setattr(provider, name, value) provider.finalize() if role == "critic": diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 0a79b99ded..46e6a8f637 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -300,6 +300,28 @@ def add_train_arguments(parser): default=False, ) + # Forwarded to the Megatron Bridge provider in bridge mode. Unlike + # --freeze-params-name-list, which masks gradients on the slime side, + # these are honored by the provider when the model is built. + parser.add_argument( + "--freeze-language-model", + action="store_true", + default=None, + help="Freeze the language model when building a model via Megatron Bridge.", + ) + parser.add_argument( + "--freeze-vision-model", + action="store_true", + default=None, + help="Freeze the vision encoder when building a model via Megatron Bridge.", + ) + parser.add_argument( + "--freeze-vision-projection", + action="store_true", + default=None, + help="Freeze the vision projection when building a model via Megatron Bridge.", + ) + return parser # rollout diff --git a/slime_plugins/megatron_bridge/__init__.py b/slime_plugins/megatron_bridge/__init__.py index a0425d491b..8116dd1332 100644 --- a/slime_plugins/megatron_bridge/__init__.py +++ b/slime_plugins/megatron_bridge/__init__.py @@ -1 +1,2 @@ import slime_plugins.megatron_bridge.glm4v_moe # noqa: F401 # register GLM-4.6V bridge +import slime_plugins.megatron_bridge.qwen3_5_vl # noqa: F401 # adapt Qwen3.5-VL GDN to packed sequences diff --git a/slime_plugins/megatron_bridge/qwen3_5_vl.py b/slime_plugins/megatron_bridge/qwen3_5_vl.py new file mode 100644 index 0000000000..af8b81553d --- /dev/null +++ b/slime_plugins/megatron_bridge/qwen3_5_vl.py @@ -0,0 +1,201 @@ +"""Qwen3.5-VL support on top of the official Megatron Bridge providers. + +Megatron Bridge ships ``Qwen35VLBridge`` / ``Qwen35VLMoEBridge`` (since v0.4.0) +and registers them on import, so no registration shim is needed here. What is +missing is the interaction with slime's data layout: the providers build their +Gated DeltaNet (GDN) layers from megatron-core's ``experimental_attention_variant += "gated_delta_net"``, and megatron-core's ``GatedDeltaNet.forward`` rejects +packed sequences outright:: + + if packed_seq_params is not None: + raise NotImplementedError("GDN does not support packed sequence for now.") + +Since slime removed BSHD support, every microbatch is packed THD, so the first +GDN layer raises and the bridge path is unusable for Qwen3.5-VL. + +With ``--micro-batch-size 1`` and no dynamic batching (already required for GDN, +see examples/geo3k_vlm/README.md) a THD microbatch holds a single sequence +optionally followed by right padding, and the hidden states are ``[T, 1, H]`` -- +exactly the layout the unpacked GDN path expects. GDN is causal, so trailing +padding cannot affect the outputs of the real tokens. We therefore subclass the +megatron-core module and drop the packed metadata after checking that we really +are in that regime. + +The subclass adds no parameters and renames nothing, so every GDN weight mapping +in the official bridge (in_proj / conv1d / A_log / dt_bias / out_norm / out_proj) +keeps working for checkpoint load, ``--save-hf`` and weight sync to SGLang. + +Not supported, and rejected loudly rather than silently miscomputed: + * context parallel -- slime hands each rank a zigzag slice while cu_seqlens + stays global, and GDN's recurrence cannot be split that way; + * micro batch size > 1 and dynamic batching -- several real sequences in one + microbatch would be fused into a single recurrent stream. + +Lifting the last restriction means forwarding cu_seqlens into the varlen kernels +instead of dropping it, which is left to a follow-up. + +Checked on Qwen3.5-2B with Megatron Bridge 0.5.0 and megatron-core 0.16.0rc0: +without this module the first GDN layer raises, with it a packed forward runs and +its logits agree with HuggingFace on 47/48 argmax positions (bf16, TP=1), with +every Megatron top-1 inside the HuggingFace top-5. Exporting back to HuggingFace +yields 621/632 tensors -- the 11 absent ones are MTP layers, which this +configuration does not build -- including all 162 GDN tensors, with a maximum +weight difference of 0.0039, i.e. one bf16 ulp. + +Gradients were checked separately at TP=2 with sequence length 512: finite +throughout, including all 90 GDN parameters. That check needs the example's +``--attention-backend flash``; letting Megatron choose picks Transformer +Engine's cuDNN fused attention, whose backward goes non-finite for this model +under bf16 with packed sequences and TP above 1. +""" + +import contextlib +import functools + +from megatron.bridge.models import gpt_provider +from megatron.bridge.models.conversion.param_mapping import AutoMapping +from megatron.bridge.models.qwen_vl import qwen35_vl_provider +from megatron.core import mpu +from megatron.core.ssm.gated_delta_net import GatedDeltaNet +from megatron.training import get_args + + +class SinglePackedSequenceGatedDeltaNet(GatedDeltaNet): + """GDN that accepts slime's THD microbatches holding one real sequence.""" + + def forward(self, hidden_states, attention_mask=None, *args, packed_seq_params=None, **kwargs): + if packed_seq_params is None: + return super().forward(hidden_states, attention_mask, *args, **kwargs) + + if mpu.get_context_parallel_world_size() != 1: + raise NotImplementedError( + "Qwen3.5-VL GDN does not support context parallel; use --context-parallel-size 1." + ) + if packed_seq_params.qkv_format != "thd": + raise NotImplementedError(f"Qwen3.5-VL GDN expects thd packing, got {packed_seq_params.qkv_format}.") + if getattr(packed_seq_params, "cu_seqlens_q_padded", None) is not None: + raise NotImplementedError("Qwen3.5-VL GDN does not support pre-padded cu_seqlens.") + if hidden_states.shape[1] != 1: + raise NotImplementedError(f"Qwen3.5-VL GDN expects a batch dimension of 1, got {hidden_states.shape[1]}.") + + total_tokens = hidden_states.shape[0] * self.sp_size + if int(packed_seq_params.cu_seqlens_q[-1]) != total_tokens: + raise RuntimeError( + f"cu_seqlens[-1]={int(packed_seq_params.cu_seqlens_q[-1])} does not match " + f"{total_tokens} tokens; the microbatch is not a single packed sequence." + ) + + return super().forward(hidden_states, attention_mask, *args, **kwargs) + + +def _patch_gated_delta_net_specs(block_spec) -> None: + """Swap GDN modules in a block spec, mirroring the bridge's own spec patching. + + Standard attention layers are left alone; only layers whose self_attention is + a megatron-core GatedDeltaNet are replaced. + """ + if block_spec is None: + return + + layer_specs = getattr(block_spec, "layer_specs", None) + if layer_specs is not None: + for layer_spec in layer_specs: + _patch_gated_delta_net_specs(layer_spec) + return + + submodules = getattr(block_spec, "submodules", None) + if submodules is None: + return + + if hasattr(submodules, "mtp_model_layer"): + _patch_gated_delta_net_specs(submodules.mtp_model_layer) + + attention_spec = getattr(submodules, "self_attention", None) + if attention_spec is None: + return + module = getattr(attention_spec, "module", None) + if isinstance(module, type) and issubclass(module, GatedDeltaNet): + attention_spec.module = SinglePackedSequenceGatedDeltaNet + + +def _check_single_sequence_microbatches() -> None: + """Reject configurations that put more than one real sequence per microbatch. + + Unlike the checks in ``forward``, this cannot be detected from the packed + metadata: a microbatch holding one sequence plus right padding and one + holding two sequences both expose two cu_seqlens segments. + """ + args = get_args() + if getattr(args, "use_dynamic_batch_size", False): + raise NotImplementedError( + "Qwen3.5-VL GDN packs several sequences per microbatch under " + "--use-dynamic-batch-size, which its recurrence cannot separate; drop the flag." + ) + if getattr(args, "micro_batch_size", 1) != 1: + raise NotImplementedError(f"Qwen3.5-VL GDN requires --micro-batch-size 1, got {args.micro_batch_size}.") + + +def _wrap_spec_builder(builder): + """Wrap the provider's block spec builder so the GDN layers it emits get swapped.""" + if getattr(builder, "_slime_patches_gdn", False): + return builder + + @functools.wraps(builder) + def build_spec(*args, **kwargs): + spec = builder(*args, **kwargs) + _patch_gated_delta_net_specs(spec) + return spec + + build_spec._slime_patches_gdn = True + return build_spec + + +@contextlib.contextmanager +def _gdn_spec_patching(): + """Wrap the spec builders for the duration of one provide() call. + + The Qwen3.5-VL providers call the block spec builders as module globals rather + than through ``self.transformer_layer_spec``, so they are wrapped in place and + restored afterwards to keep other models untouched. + """ + targets = [] + for module, name in ( + (qwen35_vl_provider, "get_transformer_block_with_experimental_attention_variant_spec"), + (gpt_provider, "mtp_block_spec"), + ): + builder = getattr(module, name, None) + if callable(builder): + targets.append((module, name, builder)) + setattr(module, name, _wrap_spec_builder(builder)) + try: + yield + finally: + for module, name, builder in targets: + setattr(module, name, builder) + + +def _wrap_provide(provider_cls, method_name: str) -> None: + original = getattr(provider_cls, method_name) + + def provide(self, *args, _original=original, **kwargs): + _check_single_sequence_microbatches() + # The language-model-only path goes through GPTModelProvider.provide, which + # does read the builder off the instance. + builder = self.transformer_layer_spec + if callable(builder): + self.transformer_layer_spec = _wrap_spec_builder(builder) + else: + _patch_gated_delta_net_specs(builder) + with _gdn_spec_patching(): + return _original(self, *args, **kwargs) + + setattr(provider_cls, method_name, provide) + + +# AutoMapping dispatches on the exact module class name, so the subclass has to be +# registered the same way megatron-bridge registers GatedDeltaNet itself. +AutoMapping.register_module_type(SinglePackedSequenceGatedDeltaNet.__name__, "column") + +for _provider_cls in (qwen35_vl_provider.Qwen35VLModelProvider, qwen35_vl_provider.Qwen35VLMoEModelProvider): + for _method_name in ("provide", "provide_language_model"): + _wrap_provide(_provider_cls, _method_name) diff --git a/tests/test_qwen3_5_vl_gdn_packed.py b/tests/test_qwen3_5_vl_gdn_packed.py new file mode 100644 index 0000000000..61537c584e --- /dev/null +++ b/tests/test_qwen3_5_vl_gdn_packed.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +import importlib.util +import pathlib +import sys +import types + +import pytest +import torch +import torch.nn as nn + +NUM_GPUS = 0 + +# Loaded straight from its path: importing it as a package member would pull in the +# other bridge plugins, which need a real megatron.bridge. +PLUGIN_PATH = pathlib.Path(__file__).resolve().parents[1] / "slime_plugins" / "megatron_bridge" / "qwen3_5_vl.py" +PLUGIN = "qwen3_5_vl_under_test" + + +class _StubGatedDeltaNet(nn.Module): + """Stands in for megatron-core's GatedDeltaNet, including its packed refusal.""" + + def __init__(self, config=None): + super().__init__() + self.config = config + self.sp_size = 1 + + def forward(self, hidden_states, attention_mask=None, *args, packed_seq_params=None, **kwargs): + if packed_seq_params is not None: + raise NotImplementedError("GDN does not support packed sequence for now.") + return hidden_states + + +class _StubSelfAttention(nn.Module): + pass + + +class _ModuleSpec: + def __init__(self, module=None, submodules=None): + self.module = module + self.submodules = submodules + + +class _Submodules: + def __init__(self, self_attention=None, mtp_model_layer=None): + self.self_attention = self_attention + if mtp_model_layer is not None: + self.mtp_model_layer = mtp_model_layer + + +class _BlockSpec: + def __init__(self, layer_specs): + self.layer_specs = layer_specs + + +def _install_stubs(monkeypatch, *, cp_size=1, micro_batch_size=1, use_dynamic_batch_size=False): + """Install the minimal megatron / megatron.bridge surface the plugin imports.""" + registered: dict[str, str] = {} + + def module(name): + mod = types.ModuleType(name) + monkeypatch.setitem(sys.modules, name, mod) + return mod + + module("megatron") + core = module("megatron.core") + module("megatron.core.ssm") + gdn_mod = module("megatron.core.ssm.gated_delta_net") + gdn_mod.GatedDeltaNet = _StubGatedDeltaNet + + mpu = types.SimpleNamespace(get_context_parallel_world_size=lambda: cp_size) + core.mpu = mpu + monkeypatch.setitem(sys.modules, "megatron.core.mpu", mpu) + + training = module("megatron.training") + training.get_args = lambda: types.SimpleNamespace( + micro_batch_size=micro_batch_size, + use_dynamic_batch_size=use_dynamic_batch_size, + ) + + module("megatron.bridge") + models_mod = module("megatron.bridge.models") + module("megatron.bridge.models.conversion") + param_mapping = module("megatron.bridge.models.conversion.param_mapping") + + class AutoMapping: + @classmethod + def register_module_type(cls, name, parallelism_type): + registered[name] = parallelism_type + + param_mapping.AutoMapping = AutoMapping + + gpt_provider = module("megatron.bridge.models.gpt_provider") + gpt_provider.mtp_block_spec = lambda config, vp_stage=None: _make_block_spec() + models_mod.gpt_provider = gpt_provider + + module("megatron.bridge.models.qwen_vl") + provider_mod = module("megatron.bridge.models.qwen_vl.qwen35_vl_provider") + provider_mod.get_transformer_block_with_experimental_attention_variant_spec = ( + lambda config, vp_stage=None: _make_block_spec() + ) + + class Qwen35VLModelProvider: + """Mimics the real provider: builds its block spec from module globals.""" + + transformer_layer_spec = staticmethod(lambda config, vp_stage=None: _make_block_spec()) + + def provide(self, pre_process=None, post_process=None, vp_stage=None): + block_spec = provider_mod.get_transformer_block_with_experimental_attention_variant_spec( + self, vp_stage=vp_stage + ) + mtp_spec = gpt_provider.mtp_block_spec(self, vp_stage=vp_stage) + return block_spec, mtp_spec + + def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None): + return self.transformer_layer_spec(self) + + class Qwen35VLMoEModelProvider(Qwen35VLModelProvider): + pass + + provider_mod.Qwen35VLModelProvider = Qwen35VLModelProvider + provider_mod.Qwen35VLMoEModelProvider = Qwen35VLMoEModelProvider + + sys.modules.pop(PLUGIN, None) + spec = importlib.util.spec_from_file_location(PLUGIN, PLUGIN_PATH) + plugin = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, PLUGIN, plugin) + spec.loader.exec_module(plugin) + return plugin, provider_mod, registered + + +def _make_block_spec(): + """One GDN layer, one standard attention layer, and a GDN layer nested under MTP.""" + gdn_layer = _ModuleSpec(submodules=_Submodules(self_attention=_ModuleSpec(module=_StubGatedDeltaNet))) + attn_layer = _ModuleSpec(submodules=_Submodules(self_attention=_ModuleSpec(module=_StubSelfAttention))) + mtp_layer = _ModuleSpec( + submodules=_Submodules( + self_attention=_ModuleSpec(module=_StubSelfAttention), + mtp_model_layer=_ModuleSpec(submodules=_Submodules(self_attention=_ModuleSpec(module=_StubGatedDeltaNet))), + ) + ) + return _BlockSpec([gdn_layer, attn_layer, mtp_layer]) + + +def _packed_seq_params(total_tokens, qkv_format="thd"): + return types.SimpleNamespace( + qkv_format=qkv_format, + cu_seqlens_q=torch.tensor([0, total_tokens], dtype=torch.int32), + cu_seqlens_q_padded=None, + ) + + +def test_only_gdn_layers_are_replaced(monkeypatch): + plugin, provider_mod, _ = _install_stubs(monkeypatch) + + spec, _mtp = provider_mod.Qwen35VLModelProvider().provide() + gdn_layer, attn_layer, mtp_layer = spec.layer_specs + + assert gdn_layer.submodules.self_attention.module is plugin.SinglePackedSequenceGatedDeltaNet + assert attn_layer.submodules.self_attention.module is _StubSelfAttention + # nested MTP layers are reached too + nested = mtp_layer.submodules.mtp_model_layer.submodules.self_attention + assert nested.module is plugin.SinglePackedSequenceGatedDeltaNet + + +def test_moe_provider_is_patched_too(monkeypatch): + plugin, provider_mod, _ = _install_stubs(monkeypatch) + + spec, _mtp = provider_mod.Qwen35VLMoEModelProvider().provide() + assert spec.layer_specs[0].submodules.self_attention.module is plugin.SinglePackedSequenceGatedDeltaNet + + +def test_subclass_is_registered_for_weight_mapping(monkeypatch): + plugin, _, registered = _install_stubs(monkeypatch) + + # AutoMapping dispatches on the exact class name, so the subclass must be registered + # the same way megatron-bridge registers GatedDeltaNet, or weight conversion fails. + assert registered[plugin.SinglePackedSequenceGatedDeltaNet.__name__] == "column" + + +def test_packed_single_sequence_is_accepted(monkeypatch): + plugin, _, _ = _install_stubs(monkeypatch) + + layer = plugin.SinglePackedSequenceGatedDeltaNet() + hidden = torch.zeros(8, 1, 4) + + # the unpatched parent would raise NotImplementedError here + out = layer(hidden, None, packed_seq_params=_packed_seq_params(8)) + assert out.shape == hidden.shape + + +def test_unpacked_input_is_delegated_unchanged(monkeypatch): + plugin, _, _ = _install_stubs(monkeypatch) + + layer = plugin.SinglePackedSequenceGatedDeltaNet() + hidden = torch.zeros(8, 2, 4) + assert layer(hidden, None).shape == hidden.shape + + +@pytest.mark.parametrize( + "kwargs, packed, message", + [ + ({"cp_size": 2}, _packed_seq_params(8), "context parallel"), + ({}, _packed_seq_params(8, qkv_format="bshd"), "thd"), + ({}, _packed_seq_params(5), "not a single packed sequence"), + ], +) +def test_unsupported_packing_is_rejected(monkeypatch, kwargs, packed, message): + plugin, _, _ = _install_stubs(monkeypatch, **kwargs) + + layer = plugin.SinglePackedSequenceGatedDeltaNet() + with pytest.raises((NotImplementedError, RuntimeError), match=message): + layer(torch.zeros(8, 1, 4), None, packed_seq_params=packed) + + +def test_batch_dimension_greater_than_one_is_rejected(monkeypatch): + plugin, _, _ = _install_stubs(monkeypatch) + + layer = plugin.SinglePackedSequenceGatedDeltaNet() + with pytest.raises(NotImplementedError, match="batch dimension"): + layer(torch.zeros(8, 2, 4), None, packed_seq_params=_packed_seq_params(8)) + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"use_dynamic_batch_size": True}, "use-dynamic-batch-size"), + ({"micro_batch_size": 4}, "micro-batch-size 1"), + ], +) +def test_multi_sequence_microbatches_are_rejected_at_build_time(monkeypatch, kwargs, message): + # These cannot be caught in forward: one sequence plus padding and two sequences + # both look like two cu_seqlens segments. + _, provider_mod, _ = _install_stubs(monkeypatch, **kwargs) + + with pytest.raises(NotImplementedError, match=message): + provider_mod.Qwen35VLModelProvider().provide() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))