From 7067657c30471fa04ba46e59001fc4fdd5df9a65 Mon Sep 17 00:00:00 2001 From: Toby Yang <1658685923yyz7@gmail.com> Date: Sat, 25 Jul 2026 04:28:47 +0000 Subject: [PATCH] [Multimodal][Model] Make Qwen3.5-VL work with packed sequences Megatron Bridge has shipped Qwen35VLBridge / Qwen35VLMoEBridge since v0.4.0 and registers them on import, so the forked Megatron Bridge pinned by examples/geo3k_vlm/run_geo3k_qwen35.sh is no longer needed. What is missing is the interaction with our data layout. The providers build their Gated DeltaNet layers from megatron-core's experimental_attention_variant="gated_delta_net", and megatron-core's GatedDeltaNet.forward raises NotImplementedError when it is given packed sequences. Since #2100 removed BSHD every microbatch is packed THD, so the first GDN layer raises and the bridge path cannot run Qwen3.5-VL at all. With --micro-batch-size 1 and no dynamic batching, which GDN already requires, a THD microbatch is a single sequence plus right padding and the hidden states are [T, 1, H] -- the layout the unpacked path expects, and GDN is causal so the trailing padding cannot affect the real tokens. Subclass the megatron-core module and drop the packed metadata once that regime is verified. The subclass adds no parameters and renames nothing, so the official GDN weight mappings still work for checkpoint load, --save-hf and weight sync to SGLang. Configurations that would put several real sequences in one microbatch, and context parallel, are rejected instead of silently fused into one recurrence. Verified on Qwen3.5-2B with Megatron Bridge 0.5.0 and megatron-core 0.16.0rc0. The unpatched model raises in the first GDN layer. The patched one: * runs a packed forward whose logits agree with HuggingFace on 47/48 argmax positions (bf16), with every Megatron top-1 inside the HuggingFace top-5, at both TP=1 and TP=2; * runs a real image forward with pixel_values and image_grid_thw; * exports back to HuggingFace as 621/632 tensors -- the 11 absent ones are MTP layers this configuration does not build -- including all 162 GDN tensors, with a maximum weight difference of one bf16 ulp; * shards the GDN weights under TP=2, which is what the AutoMapping registration controls; * produces finite gradients at TP=2 with sequence length 512, and completes an optimizer step with grad_norm 135.6. Not covered: MoE, weight sync to a live SGLang engine, and PP above 1. Also forward moe_aux_loss_coeff and the freeze_* flags to bridge providers, so the geo3k_vlm README no longer has to ask users to edit model_provider.py. --- .github/workflows/pr-test.yml | 4 + .github/workflows/pr-test.yml.j2 | 1 + examples/geo3k_vlm/README.md | 11 +- examples/geo3k_vlm/run_geo3k_qwen35.sh | 10 +- .../backends/megatron_utils/model_provider.py | 16 ++ slime/utils/arguments.py | 22 ++ slime_plugins/megatron_bridge/__init__.py | 1 + slime_plugins/megatron_bridge/qwen3_5_vl.py | 201 +++++++++++++++ tests/test_qwen3_5_vl_gdn_packed.py | 241 ++++++++++++++++++ 9 files changed, 494 insertions(+), 13 deletions(-) create mode 100644 slime_plugins/megatron_bridge/qwen3_5_vl.py create mode 100644 tests/test_qwen3_5_vl_gdn_packed.py 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__]))