Module `neuralops.tests.test_tpa_comprehensive`

Functions

def make_config(hidden_size=64, num_heads=4, num_kv_heads=2, is_decoder=False, model_type='bert')

Expand source code

def make_config(hidden_size=64, num_heads=4, num_kv_heads=2, is_decoder=False, model_type="bert"):
    return types.SimpleNamespace(
        hidden_size=hidden_size,
        num_attention_heads=num_heads,
        num_key_value_heads=num_kv_heads,
        head_dim=hidden_size // num_heads,
        is_decoder=is_decoder,
        model_type=model_type,
        dropout=0.0,
        lora_rank=16,
    )

def test_tpa_causal_mode()

Expand source code

def test_tpa_causal_mode():
    cfg = make_config(hidden_size=64, num_heads=4, num_kv_heads=2, is_decoder=True, model_type="gpt2")
    attn = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), is_decoder=True, rank=4)
    attn.eval()

    B, L, D = 2, 16, 64
    x = torch.randn(B, L, D)
    out, _ = attn(x)
    assert out.shape == (B, L, D)
    assert not torch.isnan(out).any()

def test_tpa_forward_and_shapes()

Expand source code

def test_tpa_forward_and_shapes():
    cfg = make_config(hidden_size=64, num_heads=4, num_kv_heads=2)
    attn = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4)
    attn.eval()

    for seq_len in [2, 8, 16, 25]:
        B, L, D = 2, seq_len, 64
        x = torch.randn(B, L, D)
        out, _ = attn(x)
        assert out.shape == (B, L, D), f"Expected shape {(B, L, D)}, got {out.shape}"
        assert not torch.isnan(out).any()

def test_tpa_gpt2_is_decoder_detection()

Expand source code

def test_tpa_gpt2_is_decoder_detection():
    class DummyGPT2Attention(nn.Module):
        def __init__(self, dim=64):
            super().__init__()
            self.c_attn = nn.Linear(dim, 3 * dim)
            self.c_proj = nn.Linear(dim, dim)
            self.is_decoder = True

    gpt_layer = DummyGPT2Attention(64)
    cfg = make_config(hidden_size=64, num_heads=4, is_decoder=False)
    attn = TPAttention(cfg, pretrained_attn=gpt_layer, rank=4)
    assert attn.is_decoder is True

def test_tpa_gradient_checkpointing_toggle()

Expand source code

def test_tpa_gradient_checkpointing_toggle():
    cfg = make_config(hidden_size=64, num_heads=4)
    # 1. With gradient checkpointing enabled in training
    attn_ckpt = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4, gradient_checkpointing=True)
    attn_ckpt.train()
    x1 = torch.randn(2, 8, 64, requires_grad=True)
    out1, _ = attn_ckpt(x1)
    loss1 = out1.sum()
    loss1.backward()
    assert x1.grad is not None
    assert not torch.isnan(x1.grad).any()

    # 2. With gradient checkpointing disabled
    attn_no_ckpt = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4, gradient_checkpointing=False)
    attn_no_ckpt.train()
    x2 = torch.randn(2, 8, 64, requires_grad=True)
    out2, _ = attn_no_ckpt(x2)
    loss2 = out2.sum()
    loss2.backward()
    assert x2.grad is not None
    assert not torch.isnan(x2.grad).any()

def test_tpa_inject()

Expand source code

def test_tpa_inject():
    class DummyLayer(nn.Module):
        def __init__(self, d=64, h=4):
            super().__init__()
            self.attention = DummyAttn(d, h)

    layer = DummyLayer(64, 4)
    cfg = make_config(hidden_size=64, num_heads=4, num_kv_heads=2)
    injected = inject_tpa(layer, model_config=cfg, old_attn=layer.attention)
    assert isinstance(injected, TPAttention)
    assert layer.attention is injected

    empty_layer = nn.Module()
    # assert inject_tpa(empty_layer, model_config=cfg, old_attn=None) is None

def test_tpa_lora_and_gradients()

Expand source code

def test_tpa_lora_and_gradients():
    cfg = make_config(hidden_size=64, num_heads=4, num_kv_heads=2)
    attn = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4, lora_rank=8)
    attn.train()

    # Base frozen, lora trainable
    assert not attn.v_proj.weight.requires_grad
    assert attn.v_proj.lora_A.requires_grad
    assert attn.v_proj.lora_B.requires_grad

    x = torch.randn(2, 8, 64, requires_grad=True)
    out, _ = attn(x)
    loss = out.sum()
    loss.backward()

    assert attn.v_proj.lora_A.grad is not None
    assert attn.v_proj.lora_B.grad is not None
    assert attn.v_proj.weight.grad is None

def test_tpa_masks()

Expand source code

def test_tpa_masks():
    cfg = make_config(hidden_size=64, num_heads=4, num_kv_heads=2)
    attn = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4)
    attn.eval()

    B, L, D = 2, 12, 64
    x = torch.randn(B, L, D)

    # 2D mask
    mask_2d = torch.ones(B, L)
    mask_2d[:, -3:] = 0
    out_2d, _ = attn(x, attention_mask=mask_2d)
    assert out_2d.shape == (B, L, D)
    assert not torch.isnan(out_2d).any()

    # 4D mask
    mask_4d = torch.zeros(B, 1, L, L)
    mask_4d[:, :, :, -3:] = -10000.0
    out_4d, _ = attn(x, attention_mask=mask_4d)
    assert out_4d.shape == (B, L, D)
    assert not torch.isnan(out_4d).any()

def test_tpa_nested_encoder_injection()

Expand source code

def test_tpa_nested_encoder_injection():
    class DummyBertSelf(nn.Module):
        def __init__(self):
            super().__init__()
            self.query = nn.Linear(64, 64)
            self.key = nn.Linear(64, 64)
            self.value = nn.Linear(64, 64)

    class DummyBertAttention(nn.Module):
        def __init__(self):
            super().__init__()
            self.self = DummyBertSelf()
            self.output = nn.Module()
            self.output.dense = nn.Linear(64, 64)

    class DummyBertLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = DummyBertAttention()

    class DummyLlamaLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.self_attn = DummyAttn(64, 4)

    bert_layer = DummyBertLayer()
    injected_bert = inject_tpa(bert_layer, model_config=make_config(64, 4), old_attn=bert_layer.attention.self)
    assert injected_bert.is_nested_encoder is True
    assert bert_layer.attention.self is injected_bert

    llama_layer = DummyLlamaLayer()
    injected_llama = inject_tpa(llama_layer, model_config=make_config(64, 4), old_attn=llama_layer.self_attn)
    assert injected_llama.is_nested_encoder is False
    assert llama_layer.self_attn is injected_llama

def test_tpa_no_output_scale()

Expand source code

def test_tpa_no_output_scale():
    cfg = make_config(hidden_size=64, num_heads=4)
    attn = TPAttention(cfg, pretrained_attn=DummyAttn(64, 4), rank=4)
    assert not hasattr(attn.q_proj1, "output_scale")
    assert not hasattr(attn.q_proj2, "output_scale")
    assert not hasattr(attn.v_proj, "output_scale")
    assert not hasattr(attn.out_proj, "output_scale")

def test_tpa_non_square_head_dim()

Expand source code

def test_tpa_non_square_head_dim():
    cfg = types.SimpleNamespace(
        hidden_size=512,
        num_attention_heads=4,
        num_key_value_heads=4,
        head_dim=128,
        is_decoder=False,
        model_type="bert",
        dropout=0.0,
    )
    attn = TPAttention(cfg, rank=8)
    assert attn.head_dim == 128
    assert attn.rank == 8
    assert attn.qk_proj is not None
    x = torch.randn(2, 6, 512)
    out, _ = attn(x)
    assert out.shape == (2, 6, 512)

def test_tpa_orthogonal_init()

Expand source code

def test_tpa_orthogonal_init():
    cfg = make_config(hidden_size=64, num_heads=4)
    dummy = DummyAttn(64, 4)
    attn = TPAttention(cfg, pretrained_attn=dummy, rank=4)

    w1 = attn.q_proj1.weight
    w2 = attn.q_proj2.weight
    assert float(torch.linalg.norm(w1)) > 0.5
    assert float(torch.linalg.norm(w2)) > 0.5
    assert not torch.isnan(w1).any()
    assert not torch.isnan(w2).any()

def test_tpa_qwen_head_dim_and_rank()

Expand source code

def test_tpa_qwen_head_dim_and_rank():
    class Qwen35Config:
        hidden_size = 1024
        num_attention_heads = 8
        num_key_value_heads = 8
        head_dim = 256
        model_type = "qwen3_5"
        is_decoder = True
        dropout = 0.0

    attn = TPAttention(Qwen35Config())
    assert attn.head_dim == 256
    assert attn.q_head_size == 2048
    assert attn.rank == 8  # 16*16 = 256
    assert attn.is_decoder is True
    assert attn.qk_proj is not None

    x = torch.randn(2, 4, 1024)
    out, _ = attn(x)
    assert out.shape == (2, 4, 1024)

def test_tpa_rotate_half_and_odd_dim()

Expand source code

def test_tpa_rotate_half_and_odd_dim():
    for d in [16, 17, 31, 32]:
        t = torch.randn(2, 4, 8, d)
        rotated = _rotate_half(t)
        assert rotated.shape == t.shape, f"Odd dim {d} failed in _rotate_half"

    head_dim_odd = 15
    q = torch.randn(2, 4, 8, head_dim_odd)
    k = torch.randn(2, 2, 8, head_dim_odd)
    cos = torch.randn(8, head_dim_odd)
    sin = torch.randn(8, head_dim_odd)

    q_rot, k_rot = apply_rotary_pos_emb(q, k, cos, sin)
    assert q_rot.shape == q.shape
    assert k_rot.shape == k.shape
    assert not torch.isnan(q_rot).any()

Classes

class DummyAttn (d=64, h=4)

Expand source code

class DummyAttn(nn.Module):
    def __init__(self, d=64, h=4):
        super().__init__()
        self.q_proj = nn.Linear(d, d)
        self.k_proj = nn.Linear(d, d)
        self.v_proj = nn.Linear(d, d)
        self.o_proj = nn.Linear(d, d)
        self.hidden_size = d
        self.num_heads = h

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

torch.nn.modules.module.Module