diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py index 7ddaf6c3a1b..cceaadd8e91 100644 --- a/megatron/core/ssm/gated_delta_net/common.py +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -261,8 +261,11 @@ def __init__( ) self.recompute_norm_out = False self.norm_out_checkpoint = None + self.recompute_in_proj_conv = False + # recompute_gated_delta_rule may need a checkpoint instance (same as norm_out) later if self.config.recompute_granularity == "selective": self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + self.recompute_in_proj_conv = "gdn_in_proj_conv" in self.config.recompute_modules self.out_proj = build_module( submodules.out_proj, @@ -366,6 +369,101 @@ def _gated_norm_and_a2a( return norm_out + def _in_proj_conv( + self, hidden_states, *, thd_cp_a2a_idx, batch, seq_len, cu_seqlens_q, packed_seq_params + ): + """inproj + all2all + split view + conv1d""" + # Input projection + nvtx_range_push(suffix="in_proj") + qkvzba, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + + # CP All to All: CP to HP + if self.cp_size > 1: + # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. + head_perm = _build_head_perm_for_split_sections( + self.in_proj_split_sections, + self.pg_collection.cp.size(), + torch.cuda.current_device(), + ) + qkvzba = qkvzba.index_select(-1, head_perm) + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + undo_attention_load_balancing=False, + ) + if self.cp_size > 1: + qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) + else: + qkvzba = tensor_a2a_cp2hp( + qkvzba, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + + # Transpose: s b x --> b s x + # From sbhd to bshd format + qkvzba = qkvzba.transpose(0, 1) + + # Split the tensor into q, k, v, gate (z), and the variant-specific gate features + # (beta, alpha for GDN; f, b, w for GDN2) + qkv, gate, beta, alpha = torch.split(qkvzba, self.feat_dim_split, dim=-1) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + + # Convolution on qkv + nvtx_range_push(suffix="conv1d") + seq_len = qkv.shape[1] + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=self.pg_collection.cp, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + if self.config.deterministic_mode: + qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s + conv_out = F.conv1d( + input=qkv, # Torch-native only accept [b, d, s] format input + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim_local_tp // self.cp_size, + ) + qkv = self.act_fn(conv_out[..., :seq_len]) + qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d + else: + assert self.activation in ["silu", "swish"] + qkv, _ = causal_conv1d( + x=qkv, # FLA conv1d accepts [b, s, d] format input + weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens_q, + ) + nvtx_range_pop(suffix="conv1d") + + return qkv, gate, beta, alpha + @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py index 65d9dc7df0a..9432bd90c80 100644 --- a/megatron/core/ssm/gated_delta_net/gdn.py +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -9,19 +9,15 @@ from typing import Optional import torch -import torch.nn.functional as F from megatron.core import tensor_parallel from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.ssm.gated_delta_net.common import ( - _build_head_perm_for_split_sections, _build_thd_cp_a2a_perm, _GDNBase, - causal_conv1d, chunk_gated_delta_rule, get_parameter_local_cp, - tensor_a2a_cp2hp, torch_chunk_gated_delta_rule, ) from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push @@ -126,30 +122,8 @@ def forward( cu_seqlens_q = None cu_seqlens_kv = None - # Input projection - nvtx_range_push(suffix="in_proj") - qkvzba, _ = self.in_proj(hidden_states) - nvtx_range_pop(suffix="in_proj") - - # CP All to All: CP to HP - if self.cp_size > 1: - # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. - head_perm = _build_head_perm_for_split_sections( - self.in_proj_split_sections, - self.pg_collection.cp.size(), - torch.cuda.current_device(), - ) - qkvzba = qkvzba.index_select(-1, head_perm) - thd_cp_a2a_idx, thd_cp_a2a_inv = None, None if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - qkvzba = tensor_a2a_cp2hp( - qkvzba, - seq_dim=0, - head_dim=-1, - cp_group=self.pg_collection.cp, - undo_attention_load_balancing=False, - ) if self.cp_size > 1: # Permute at the seq dim so that a single unsectioned a2a # is equivalent to per-sequence a2a. @@ -157,70 +131,22 @@ def forward( thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm( cu_seqlens_q, self.cp_size, seq_len ) - qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) - else: - qkvzba = tensor_a2a_cp2hp( - qkvzba, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - - # Transpose: s b x --> b s x - # From sbhd to bshd format - qkvzba = qkvzba.transpose(0, 1) - # Split the tensor into q, k, v, gate (z), and the variant-specific gate features - # (beta, alpha for GDN; f, b, w for GDN2) - qkv, gate, beta, alpha = torch.split(qkvzba, self.feat_dim_split, dim=-1) - gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - - # Convolution on qkv - nvtx_range_push(suffix="conv1d") - seq_len = qkv.shape[1] - qkv_channels_split_sections = [ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - ] - conv1d_weight = get_parameter_local_cp( - self.conv1d.weight, - dim=0, - cp_group=self.pg_collection.cp, - split_sections=qkv_channels_split_sections, - ) - conv1d_bias = ( - get_parameter_local_cp( - self.conv1d.bias, - dim=0, - cp_group=self.pg_collection.cp, - split_sections=qkv_channels_split_sections, - ) - if self.conv_bias - else None + in_proj_conv_func = partial( + self._in_proj_conv, + thd_cp_a2a_idx=thd_cp_a2a_idx, + batch=batch, + seq_len=seq_len, + cu_seqlens_q=cu_seqlens_q, + packed_seq_params=packed_seq_params, ) - if self.config.deterministic_mode: - qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s - conv_out = F.conv1d( - input=qkv, # Torch-native only accept [b, d, s] format input - weight=conv1d_weight, - bias=conv1d_bias, - stride=self.conv1d.stride, - padding=self.conv1d.padding, - dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // self.cp_size, + + if self.recompute_in_proj_conv: + qkv, gate, beta, alpha = tensor_parallel.checkpoint( + in_proj_conv_func, False, hidden_states ) - qkv = self.act_fn(conv_out[..., :seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d else: - assert self.activation in ["silu", "swish"] - qkv, _ = causal_conv1d( - x=qkv, # FLA conv1d accepts [b, s, d] format input - weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w - bias=conv1d_bias, - activation=self.activation, - initial_state=None, - output_final_state=False, - cu_seqlens=cu_seqlens_q, - ) - nvtx_range_pop(suffix="conv1d") + qkv, gate, beta, alpha = in_proj_conv_func(hidden_states) A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=self.pg_collection.cp) dt_bias_local_cp = get_parameter_local_cp( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 66cfaded213..6dcdb5d9329 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -556,6 +556,8 @@ class TransformerConfig(ModelParallelConfig): "moe": recompute the MoE layer. "shared_experts": recompute the shared experts in the MoE layer. "gdn_norm_out": recompute the GatedDeltaNet output norm and HP-to-CP all-to-all. + "gdn_in_proj_conv": recompute the GatedDeltaNet input projection, HP-to-CP all-to-all, split, + and convolution. "moe_act", "layernorm", "mla_up_proj", and "gdn_norm_out" use output-discarding checkpointing, "core_attn", "mlp", "moe", and "shared_experts" use normal checkpointing. """ @@ -1800,6 +1802,7 @@ def __post_init__(self): "moe", "shared_experts", "gdn_norm_out", + "gdn_in_proj_conv", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -1827,6 +1830,15 @@ def __post_init__(self): "experimental_attention_variant='gated_delta_net'." ) + if ( + "gdn_in_proj_conv" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_in_proj_conv in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 7cd2eb5e104..78f72093443 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -23,6 +23,7 @@ ) from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig +from megatron.core.utils import is_te_min_version from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness from tests.unit_tests.transformer.test_multi_latent_attention import ( @@ -452,6 +453,303 @@ def test_gpu_forward_thd_padding_correctness(self): with pytest.raises(ValueError, match="does not match"): self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + def test_selective_recompute_in_proj_conv_deterministic(self): + """Deterministic mode: recompute of the in_proj+conv region is bit-exact + vs no-recompute. Hard gate on all architectures.""" + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + + def build_gdn(config): + submod = get_experimental_attention_variant_module_spec(config=config).submodules + gdn = GatedDeltaNet( + config, + submodules=submod, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hs): + gdn.zero_grad(set_to_none=True) + hs.grad = None + output, _ = gdn(hs, None) + output.float().sum().backward() + grads = { + name: p.grad.detach().clone() + for name, p in gdn.named_parameters() + if p.grad is not None + } + return output.detach().clone(), grads, hs.grad.detach().clone() + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + base_config.deterministic_mode = True + rec_config = copy.deepcopy(self.transformer_config) + rec_config.deterministic_mode = True + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_in_proj_conv"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_in_proj_conv is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + del base_gdn + torch.cuda.empty_cache() + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_in_proj_conv is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + + def test_selective_recompute_in_proj_conv_within_kernel_noise(self): + """Non-deterministic mode: recompute not add error beyond the kernels' + own run-to-run noise. Tolerance is derived per-run from the baseline's + measured noise floor, so it stays valid across architectures. The printed + deviations double as the per-HW kernel-determinism characterization. + """ + pg_collection = ProcessGroupCollection( + tp=parallel_state.get_tensor_model_parallel_group(), + cp=parallel_state.get_context_parallel_group(), + ) + + def build_gdn(config): + submod = get_experimental_attention_variant_module_spec(config=config).submodules + gdn = GatedDeltaNet( + config, + submodules=submod, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hs): + gdn.zero_grad(set_to_none=True) + hs.grad = None + output, _ = gdn(hs, None) + output.float().sum().backward() + grads = { + name: p.grad.detach().clone() + for name, p in gdn.named_parameters() + if p.grad is not None + } + return output.detach().clone(), grads, hs.grad.detach().clone() + + def max_abs(a, b): + return (a - b).abs().max().item() + + micro_batch_size = 2 + seq_length = 64 + base_config = copy.deepcopy(self.transformer_config) + base_config.deterministic_mode = False + rec_config = copy.deepcopy(self.transformer_config) + rec_config.deterministic_mode = False + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_in_proj_conv"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline: run N times (NO reseed between runs) to measure the + # kernels' inherent run-to-run noise floor. --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_in_proj_conv is False + n_runs = 5 + base_runs = [run(base_gdn, hidden_states) for _ in range(n_runs)] + del base_gdn + torch.cuda.empty_cache() + + def pairwise_noise(select): + return max( + max_abs(select(base_runs[i]), select(base_runs[j])) + for i in range(n_runs) + for j in range(i + 1, n_runs) + ) + + out_noise = pairwise_noise(lambda r: r[0]) + ig_noise = pairwise_noise(lambda r: r[2]) + grad_noise = {name: pairwise_noise(lambda r, n=name: r[1][n]) for name in base_runs[0][1]} + + # --- Recompute: one run; must fall within the baseline noise floor. --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_in_proj_conv is True + rec_out, rec_grads, rec_ig = run(rec_gdn, hidden_states) + + base_out, base_grads, base_ig = base_runs[0] + out_err = max_abs(rec_out, base_out) + ig_err = max_abs(rec_ig, base_ig) + + rank = torch.distributed.get_rank() + if rank == 0: + print( + f"[in_proj_conv non-det] output: err={out_err:.3e} noise={out_noise:.3e} | " + f"input_grad: err={ig_err:.3e} noise={ig_noise:.3e}" + ) + + assert ( + out_err <= out_noise + ), f"Recompute output exceeds kernel noise: err={out_err:.3e} > noise={out_noise:.3e} ({rank=})" + assert ( + ig_err <= ig_noise + ), f"Recompute input-grad exceeds kernel noise: err={ig_err:.3e} > noise={ig_noise:.3e} ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + g_err = max_abs(rec_grads[name], base_grads[name]) + assert g_err <= grad_noise[name], ( + f"Recompute grad '{name}' exceeds kernel noise: " + f"err={g_err:.3e} > noise={grad_noise[name]:.3e} ({rank=})" + ) + + def test_in_proj_recompute_with_delay_wgrad(self): + """gdn_in_proj recompute must be compatible with delayed weight-grad compute. + + Design-doc ยง4.1: in_proj recompute reruns the linear's forward during + backward, while delay_wgrad_compute defers its weight-grad to an explicit + backward_dw() call. Both branches enable delay_wgrad_compute; the only + difference is the recompute switch, so any mismatch isolates a + recompute-vs-delayed-wgrad interaction (not TE's delay-vs-no-delay numerics). + Deterministic mode makes the comparison bitwise. + """ + if not is_te_min_version("2.3.0"): + pytest.skip("delay_wgrad_compute requires TransformerEngine >= 2.3.0") + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + def build_gdn(config): + gdn_submodules = get_experimental_attention_variant_module_spec( + config=config + ).submodules + gdn = GatedDeltaNet( + config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + return gdn.cuda().bfloat16() + + def run(gdn, hidden_states): + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + # Flush deferred weight grads (no-op when delay_wgrad_compute is off). + gdn.backward_dw() + grads = { + name: param.grad.detach() + for name, param in gdn.named_parameters() + if param.grad is not None + } + input_grad = hidden_states.grad.detach().clone() + return output.detach(), grads, input_grad + + micro_batch_size = 2 + seq_length = 64 + + # Both branches: deterministic + delayed wgrad. Only recompute differs. + base_config = copy.deepcopy(self.transformer_config) + base_config.deterministic_mode = True + base_config.delay_wgrad_compute = True + + rec_config = copy.deepcopy(self.transformer_config) + rec_config.deterministic_mode = True + rec_config.delay_wgrad_compute = True + rec_config.recompute_granularity = "selective" + rec_config.recompute_modules = ["gdn_in_proj"] + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + hidden_states = torch.randn( + ( + seq_length // self.sp_size // self.cp_size, + micro_batch_size, + self.gdn.config.hidden_size, + ), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + requires_grad=True, + ) + + # --- Baseline: delayed wgrad, NO recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + base_gdn = build_gdn(base_config) + assert base_gdn.recompute_in_proj is False + base_output, base_grads, base_input_grad = run(base_gdn, hidden_states) + hidden_states.grad = None + del base_gdn + torch.cuda.empty_cache() + + # --- Delayed wgrad + in_proj recompute --- + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + rec_gdn = build_gdn(rec_config) + assert rec_gdn.recompute_in_proj is True + rec_output, rec_grads, rec_input_grad = run(rec_gdn, hidden_states) + + rank = torch.distributed.get_rank() + assert torch.equal(rec_output, base_output), f"Output not identical ({rank=})" + assert torch.equal(rec_input_grad, base_input_grad), f"Input grad not identical ({rank=})" + assert set(rec_grads.keys()) == set(base_grads.keys()) + for name in base_grads: + assert torch.equal( + rec_grads[name], base_grads[name] + ), f"Grad not identical for {name} ({rank=})" + @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.internal @@ -723,3 +1021,80 @@ def test_cp2hp_hp2cp_round_trip(self, cu_seqlens): back = self._batched_a2a_hp2cp(mid, cu, self.cp_group) assert torch.equal(back, local_t), "Batched cp2hp -> hp2cp not identity" + + +# -----gdn selective-recompute config validation tests, cpu only----- +class TestGDNSelectiveRecomputeConfigValidation: + """selective-recompute config validation tests, cpu only.""" + + GDN_RECOMPUTE_MODULES = ["gdn_in_proj", "gdn_conv1d", "gdn_gated_delta_rule"] + + @staticmethod + def _gdn_recompute_config_kwargs(**overrides): + """Minimal kwargs for a valid gated_delta_net TransformerConfig. + + Only the fields required by ``__post_init__`` for the gated_delta_net + branch are set; everything else keeps its default. No CUDA / dist state + is touched, so this constructs fine on CPU. + """ + kwargs = dict( + num_layers=1, + hidden_size=2048, + num_attention_heads=16, + num_query_groups=2, + activation_func=F.silu, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + recompute_granularity="selective", + ) + kwargs.update(overrides) + return kwargs + + @pytest.mark.parametrize("module", GDN_RECOMPUTE_MODULES) + def test_gdn_recompute_module_accepted(self, module): + """Each new GDN module is accepted with experimental_attention_variant=gdn.""" + config = TransformerConfig(**self._gdn_recompute_config_kwargs(recompute_modules=[module])) + assert module in config.recompute_modules + + def test_gdn_recompute_all_three_accepted(self): + """All three switches together are accepted.""" + config = TransformerConfig( + **self._gdn_recompute_config_kwargs(recompute_modules=list(self.GDN_RECOMPUTE_MODULES)) + ) + assert set(self.GDN_RECOMPUTE_MODULES).issubset(set(config.recompute_modules)) + + def test_gdn_recompute_coexists_with_norm_out(self): + """New switches coexist with the pre-existing gdn_norm_out switch.""" + modules = ["gdn_norm_out", *self.GDN_RECOMPUTE_MODULES] + config = TransformerConfig(**self._gdn_recompute_config_kwargs(recompute_modules=modules)) + assert set(modules).issubset(set(config.recompute_modules)) + + @pytest.mark.parametrize("module", GDN_RECOMPUTE_MODULES) + def test_gdn_recompute_requires_gated_delta_net(self, module): + """Enabling a GDN switch without the gated_delta_net variant raises.""" + with pytest.raises( + ValueError, match="only supported with experimental_attention_variant='gated_delta_net'" + ): + TransformerConfig( + **self._gdn_recompute_config_kwargs( + recompute_modules=[module], + experimental_attention_variant=None, + # Drop gdn-only required fields so the config is otherwise valid. + linear_attention_freq=None, + linear_conv_kernel_dim=None, + linear_key_head_dim=None, + linear_value_head_dim=None, + linear_num_key_heads=None, + linear_num_value_heads=None, + ) + ) + + def test_gdn_recompute_rejects_typo(self): + """A misspelled module name is caught by the allowed_modules assert.""" + with pytest.raises(AssertionError, match="Invalid choices for recompute_modules"): + TransformerConfig(**self._gdn_recompute_config_kwargs(recompute_modules=["gdn_inproj"]))