diff --git a/tests/utils/test_megatron_gdn_patch_on_cpu.py b/tests/utils/test_megatron_gdn_patch_on_cpu.py new file mode 100644 index 00000000000..1378c0ca95a --- /dev/null +++ b/tests/utils/test_megatron_gdn_patch_on_cpu.py @@ -0,0 +1,249 @@ +# Copyright 2024 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. + +import sys +import types +from types import SimpleNamespace + +import pytest + +try: + import torch +except ModuleNotFoundError: + torch = None + + +def _install_fake_gdn_dependencies(monkeypatch, conv_backend="fla"): + """Install fake megatron/fla modules so the GDN patch can run on CPU. + + ``conv_backend`` selects which convolution path the patch will exercise: + - "fla": expose ``fla.modules.convolution.causal_conv1d`` (preferred path). + - "ccfn": omit the fla conv and expose ``causal_conv1d.causal_conv1d_fn`` + so the varlen ``seq_idx`` fallback path is taken instead. + """ + calls = {} + + class FakeGatedDeltaNet: + pass + + def torch_chunk_gated_delta_rule(query, key, value, **kwargs): + calls["torch_chunk"] = kwargs + return value, None + + def chunk_gated_delta_rule(query, key, value, **kwargs): + calls["chunk"] = kwargs + return value, None + + def fla_causal_conv1d(x, weight, bias, activation, initial_state, output_final_state, cu_seqlens): + calls["conv"] = { + "x_shape": tuple(x.shape), + "cu_seqlens": cu_seqlens, + "activation": activation, + "initial_state": initial_state, + "output_final_state": output_final_state, + } + return x, None + + def causal_conv1d_fn(x, weight, bias, activation, seq_idx): + calls["ccfn"] = { + "x_shape": tuple(x.shape), + "seq_idx": seq_idx, + "activation": activation, + } + return x + + modules = { + "megatron": types.ModuleType("megatron"), + "megatron.core": types.ModuleType("megatron.core"), + "megatron.core.ssm": types.ModuleType("megatron.core.ssm"), + "megatron.core.ssm.gated_delta_net": types.ModuleType("megatron.core.ssm.gated_delta_net"), + "megatron.core.utils": types.ModuleType("megatron.core.utils"), + "fla": types.ModuleType("fla"), + "fla.modules": types.ModuleType("fla.modules"), + "fla.modules.convolution": types.ModuleType("fla.modules.convolution"), + "fla.modules.l2norm": types.ModuleType("fla.modules.l2norm"), + "fla.ops": types.ModuleType("fla.ops"), + "fla.ops.gated_delta_rule": types.ModuleType("fla.ops.gated_delta_rule"), + } + + modules["megatron.core.ssm.gated_delta_net"].GatedDeltaNet = FakeGatedDeltaNet + modules["megatron.core.ssm.gated_delta_net"].torch_chunk_gated_delta_rule = torch_chunk_gated_delta_rule + modules["megatron.core.utils"].deprecate_inference_params = lambda inference_context, inference_params: ( + inference_context if inference_context is not None else inference_params + ) + modules["megatron.core.utils"].nvtx_range_push = lambda suffix: None + modules["megatron.core.utils"].nvtx_range_pop = lambda suffix: None + modules["fla.modules.l2norm"].l2norm = lambda x: x + modules["fla.ops.gated_delta_rule"].chunk_gated_delta_rule = chunk_gated_delta_rule + + if conv_backend == "fla": + modules["fla.modules.convolution"].causal_conv1d = fla_causal_conv1d + elif conv_backend == "ccfn": + # No fla conv: patch falls back to causal_conv1d_fn (seq_idx path). + modules["causal_conv1d"] = types.ModuleType("causal_conv1d") + modules["causal_conv1d"].causal_conv1d_fn = causal_conv1d_fn + else: + raise ValueError(conv_backend) + + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + return FakeGatedDeltaNet, calls + + +class _FakeLinear: + def __call__(self, hidden_states): + full_seq_len = hidden_states.shape[0] * 2 + batch = hidden_states.shape[1] + return torch.arange(full_seq_len * batch * 10, dtype=torch.float32).reshape(full_seq_len, batch, 10), None + + +class _FakeOutProj: + def __call__(self, norm_out): + return norm_out, None + + +class _FakeConv1d: + def __init__(self): + self.weight = torch.ones(6, 1, 3) + self.bias = torch.zeros(6) + + def __call__(self, qkv): + return qkv + + +class _FakeGdnInstance: + def __init__(self): + self.sp_size = 2 + self.config = SimpleNamespace(sequence_parallel=True, deterministic_mode=False) + self.in_proj = _FakeLinear() + self.out_proj = _FakeOutProj() + self.conv1d = _FakeConv1d() + self.activation = "silu" + self.act_fn = lambda x: x + self.qk_dim = 2 + self.v_dim = 2 + self.tp_size = 1 + self.num_value_heads = 1 + self.num_key_heads = 1 + self.value_head_dim = 2 + self.key_head_dim = 2 + self.use_qk_l2norm = True + self.A_log = torch.zeros(1) + self.dt_bias = torch.zeros(1) + + def _apply_gated_norm(self, core_attn_out, gate): + assert core_attn_out.shape == gate.shape + return core_attn_out + + +def test_gdn_patch_passes_cu_seqlens_to_fla_varlen_paths(monkeypatch): + if torch is None: + pytest.skip("torch is not installed") + + # Import before installing fake modules: importing verl.models.mcore runs its + # package __init__, which needs the real megatron.core (e.g. ModelParallelConfig). + # The patch's own megatron/fla imports are lazy, so the fakes still apply at call time. + from verl.models.mcore.patch import apply_patch_megatron_gated_delta_net + + fake_gdn_cls, calls = _install_fake_gdn_dependencies(monkeypatch) + + apply_patch_megatron_gated_delta_net() + + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + packed_seq_params = SimpleNamespace( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=None, + seq_idx=None, + ) + + hidden_states = torch.zeros(2, 1, 4) + out, bias = fake_gdn_cls.forward( + _FakeGdnInstance(), + hidden_states, + attention_mask=None, + packed_seq_params=packed_seq_params, + ) + + assert bias is None + assert tuple(out.shape) == (4, 1, 2) + assert calls["conv"]["cu_seqlens"] is cu_seqlens + assert calls["conv"]["x_shape"] == (1, 4, 6) + assert calls["chunk"]["cu_seqlens"] is cu_seqlens + + +def test_gdn_patch_builds_seq_idx_for_causal_conv1d_fn(monkeypatch): + if torch is None: + pytest.skip("torch is not installed") + + from verl.models.mcore.patch import apply_patch_megatron_gated_delta_net + + # No fla conv available -> patch falls back to causal_conv1d_fn (seq_idx path). + fake_gdn_cls, calls = _install_fake_gdn_dependencies(monkeypatch, conv_backend="ccfn") + + apply_patch_megatron_gated_delta_net() + + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + packed_seq_params = SimpleNamespace( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=None, + seq_idx=None, + ) + + hidden_states = torch.zeros(2, 1, 4) # sp_size=2 -> full seq len 4 + out, bias = fake_gdn_cls.forward( + _FakeGdnInstance(), + hidden_states, + attention_mask=None, + packed_seq_params=packed_seq_params, + ) + + assert bias is None + assert tuple(out.shape) == (4, 1, 2) + # causal_conv1d_fn receives the channel-first (b, d, s) layout. + assert calls["ccfn"]["x_shape"] == (1, 6, 4) + # seq_idx is derived from cu_seqlens: two sequences of length 2 -> [[0,0,1,1]]. + seq_idx = calls["ccfn"]["seq_idx"] + assert seq_idx is not None + assert seq_idx.dtype == torch.int32 + assert seq_idx.tolist() == [[0, 0, 1, 1]] + # cu_seqlens still flows into the varlen gated-delta-rule kernel. + assert calls["chunk"]["cu_seqlens"] is cu_seqlens + + +def test_gdn_patch_rejects_packed_deterministic_mode(monkeypatch): + if torch is None: + pytest.skip("torch is not installed") + + from verl.models.mcore.patch import apply_patch_megatron_gated_delta_net + + fake_gdn_cls, _ = _install_fake_gdn_dependencies(monkeypatch) + + apply_patch_megatron_gated_delta_net() + + instance = _FakeGdnInstance() + instance.config.deterministic_mode = True + cu_seqlens = torch.tensor([0, 2, 4], dtype=torch.int32) + packed_seq_params = SimpleNamespace( + cu_seqlens_q=cu_seqlens, + cu_seqlens_q_padded=None, + seq_idx=None, + ) + + try: + fake_gdn_cls.forward(instance, torch.zeros(2, 1, 4), attention_mask=None, packed_seq_params=packed_seq_params) + except NotImplementedError as exc: + assert "deterministic mode" in str(exc) + else: + raise AssertionError("GDN packed deterministic mode should be rejected") diff --git a/verl/models/mcore/patch.py b/verl/models/mcore/patch.py index 3e00290fa93..f1634f01248 100644 --- a/verl/models/mcore/patch.py +++ b/verl/models/mcore/patch.py @@ -571,3 +571,223 @@ def patch_backward(ctx, *args): return (None, None) + grads rd.CheckpointFunction.backward = patch_backward + + +def apply_patch_megatron_gated_delta_net(): + """Enable Megatron GatedDeltaNet for packed THD sequence-parallel inputs. + + Megatron's GatedDeltaNet currently rejects ``packed_seq_params``. verl's + packed THD path represents all valid tokens as a packed batch and relies on + ``cu_seqlens`` to preserve sequence boundaries. This patch forwards those + boundaries to the FLA gated-delta-rule implementation and to a varlen-aware + causal convolution path. + """ + try: + from megatron.core.ssm.gated_delta_net import GatedDeltaNet, torch_chunk_gated_delta_rule + except ImportError: + return + + import torch + import torch.nn.functional as F + from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push + + try: + from fla.modules.l2norm import l2norm + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + except ImportError: + return + + try: + from fla.modules.convolution import causal_conv1d as fla_causal_conv1d + except ImportError: + fla_causal_conv1d = None + + try: + from causal_conv1d import causal_conv1d_fn + except ImportError: + causal_conv1d_fn = None + + def _get_cu_seqlens(packed_seq_params): + if packed_seq_params is None: + return None + cu_seqlens_padded = getattr(packed_seq_params, "cu_seqlens_q_padded", None) + if cu_seqlens_padded is not None: + return cu_seqlens_padded + return packed_seq_params.cu_seqlens_q + + def _build_seq_idx_from_cu_seqlens(cu_seqlens, total_tokens): + if cu_seqlens is None: + return None + cu_seqlens = torch.cat([cu_seqlens, cu_seqlens.new_tensor([total_tokens])]) + seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + return ( + torch.repeat_interleave(torch.arange(seq_lengths.numel(), device=cu_seqlens.device), seq_lengths)[ + :total_tokens + ] + .to(torch.int32) + .unsqueeze(0) + ) + + def patched_forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + attention_bias=None, + packed_seq_params=None, + sequence_len_offset=None, + *, + inference_params=None, + **kwargs, + ): + inference_context = deprecate_inference_params(inference_context, inference_params) + + seq_len, batch, _ = hidden_states.shape + seq_len = seq_len * self.sp_size + + if inference_context is not None: + assert inference_context.is_static_batching(), "GDN does not currently support dynamic inference batching." + assert not self.config.sequence_parallel + raise NotImplementedError("GDN does not support inference for now.") + + cu_seqlens = _get_cu_seqlens(packed_seq_params) + seq_idx = getattr(packed_seq_params, "seq_idx", None) if packed_seq_params is not None else None + if seq_idx is None and cu_seqlens is not None: + seq_idx = _build_seq_idx_from_cu_seqlens(cu_seqlens, seq_len * batch) + + nvtx_range_push(suffix="in_proj") + qkvzba, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + + qkvzba = qkvzba.transpose(0, 1) + + qkv, gate, beta, alpha = torch.split( + qkvzba, + [ + (self.qk_dim * 2 + self.v_dim) // self.tp_size, + self.v_dim // self.tp_size, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + dim=-1, + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + beta = beta.reshape(batch, seq_len, -1) + alpha = alpha.reshape(batch, seq_len, -1) + + nvtx_range_push(suffix="conv1d") + if fla_causal_conv1d is not None and not self.config.deterministic_mode: + assert self.activation in ["silu", "swish"] + qkv = qkv.contiguous() + qkv, _ = fla_causal_conv1d( + x=qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + elif causal_conv1d_fn is not None and not self.config.deterministic_mode: + if seq_idx is None: + qkv = qkv.transpose(1, 2).contiguous() + else: + # causal_conv1d_fn's varlen path requires the channel dimension + # to be unit-stride and the outer strides to satisfy kernel + # alignment constraints. The fused projection slice is not + # guaranteed to have that layout. + qkv = qkv.contiguous().transpose(1, 2) + assert self.activation in ["silu", "swish"] + qkv = causal_conv1d_fn( + x=qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + seq_idx=seq_idx, + ) + qkv = qkv.transpose(1, 2) + else: + if seq_idx is not None or cu_seqlens is not None: + raise NotImplementedError( + "GDN packed sequence requires fla.modules.convolution.causal_conv1d " + "or causal_conv1d_fn in non-deterministic mode." + ) + qkv = qkv.transpose(1, 2).contiguous() + qkv = self.act_fn(self.conv1d(qkv)[..., :seq_len]) + qkv = qkv.transpose(1, 2) + nvtx_range_pop(suffix="conv1d") + + query, key, value = torch.split( + qkv, + [ + self.qk_dim // self.tp_size, + self.qk_dim // self.tp_size, + self.v_dim // self.tp_size, + ], + dim=-1, + ) + query = query.reshape(batch, seq_len, -1, self.key_head_dim) + key = key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + if self.use_qk_l2norm: + query = l2norm(query.contiguous()) + key = l2norm(key.contiguous()) + if self.num_value_heads // self.num_key_heads > 1: + query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + nvtx_range_push(suffix="g_and_beta") + g = -self.A_log.exp() * F.softplus(alpha.float() + self.dt_bias) + beta = beta.sigmoid() + nvtx_range_pop(suffix="g_and_beta") + + nvtx_range_push(suffix="gated_delta_rule") + if self.config.deterministic_mode: + if cu_seqlens is not None: + raise NotImplementedError("GDN packed sequence is not supported in deterministic mode.") + core_attn_out, last_recurrent_state = torch_chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + ) + else: + core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + return out, out_bias + + GatedDeltaNet.forward = patched_forward diff --git a/verl/workers/engine/megatron/transformer_impl.py b/verl/workers/engine/megatron/transformer_impl.py index 93e98f234d0..9c8c08221fd 100644 --- a/verl/workers/engine/megatron/transformer_impl.py +++ b/verl/workers/engine/megatron/transformer_impl.py @@ -130,9 +130,13 @@ def __init__( apply_mtp_inference_patch() if is_cuda_available: - from verl.models.mcore.patch import apply_patch_megatron_recomputation_backward + from verl.models.mcore.patch import ( + apply_patch_megatron_gated_delta_net, + apply_patch_megatron_recomputation_backward, + ) apply_patch_megatron_recomputation_backward() + apply_patch_megatron_gated_delta_net() def _init_device_mesh(self): # TODO: set different parallelism for actor, critic, ref