Module neuralops.tests.test_aegis_spectron_fp8

Hardening test suite for aegis_spectron_fp8.py (AegisSpectronFP8Attention).

Functions

def test_config_none()

Expand source code

def test_config_none():
    attn = AegisSpectronFP8Attention(config=None, hidden_size=64, num_heads=4, rank=16, lora_rank=8)
    x = torch.randn(1, 8, 64)
    out, _ = attn(x)
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_decoder()

Expand source code

def test_decoder():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    attn = AegisSpectronFP8Attention(config=cfg)
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_distilbert_copy()

Expand source code

def test_distilbert_copy():
    cfg = MockConfig()
    old_attn = DummyDistilBertAttention(64)
    attn = AegisSpectronFP8Attention(config=cfg, pretrained_attn=old_attn)
    assert not attn.q_proj.lin1.weight.requires_grad
    assert not attn.q_proj.lin2.weight.requires_grad
    assert attn.q_proj.lora_A.requires_grad
    assert attn.q_proj.lora_B.requires_grad
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_fail_loud()

Expand source code

def test_fail_loud():
    attn = AegisSpectronFP8Attention(config=MockConfig())
    with pytest.raises(RuntimeError, match="mandatory projections"):
        attn.copy_weights(DummyEmptyAttention())

def test_forward_stable()

Expand source code

def test_forward_stable():
    attn = AegisSpectronFP8Attention(config=MockConfig())
    out, _ = attn(torch.randn(2, 12, 64))
    assert not torch.isnan(out).any()
    assert not torch.isinf(out).any()

def test_gpt2_copy()

Expand source code

def test_gpt2_copy():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    old_attn = DummyGPT2Attention(64)
    attn = AegisSpectronFP8Attention(config=cfg, pretrained_attn=old_attn)
    assert attn.is_decoder
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_inject_distilbert()

Expand source code

def test_inject_distilbert():
    class DummyLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = DummyDistilBertAttention(64)
    layer = DummyLayer()
    res = inject_aegis_spectron_fp8(layer=layer, model_config=MockConfig())
    assert res is not None
    assert isinstance(layer.attention, AegisSpectronFP8Attention)
    out, _ = layer.attention(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_inject_none()

Expand source code

def test_inject_none():
    assert inject_aegis_spectron_fp8(layer=None, model=None) is None

def test_inject_none_on_failure()

Expand source code

def test_inject_none_on_failure():
    class DummyLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = DummyEmptyAttention()
    assert inject_aegis_spectron_fp8(layer=DummyLayer(), model_config=MockConfig()) is None

def test_qwen_gqa()

Expand source code

def test_qwen_gqa():
    cfg = MockConfig(hidden_size=64, num_attention_heads=4, num_key_value_heads=2, head_dim=16, is_decoder=True, model_type="qwen2")
    old_attn = DummyQwenAttention(dim=64, num_kv_heads=2, head_dim=16)
    attn = AegisSpectronFP8Attention(config=cfg, pretrained_attn=old_attn)
    assert attn.kv_repeat == 2
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_registry()

Expand source code

def test_registry():
    assert ATTENTION_REGISTRY.get("aegis_spectron_fp8") is not None
    assert ATTENTION_REGISTRY.get("aegis_spectron_fp8_attn") is not None

def test_roberta_copy()

Expand source code

def test_roberta_copy():
    cfg = MockConfig()
    old_attn = DummyRobertaAttention(64)
    attn = AegisSpectronFP8Attention(config=cfg, pretrained_attn=old_attn)
    assert attn.is_nested_encoder
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_rope_even_head_dim()

Expand source code

def test_rope_even_head_dim():
    attn = AegisSpectronFP8Attention(config=MockConfig())
    x = torch.randn(1, 8, 64)
    cos = torch.randn(1, 8, 16)
    sin = torch.randn(1, 8, 16)
    out, _ = attn(x, position_embeddings=(cos, sin))
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_rope_odd_head_dim()

Expand source code

def test_rope_odd_head_dim():
    cfg = MockConfig(hidden_size=68, num_attention_heads=4, head_dim=17)
    attn = AegisSpectronFP8Attention(config=cfg)
    x = torch.randn(1, 4, 68)
    cos = torch.randn(1, 4, 17)
    sin = torch.randn(1, 4, 17)
    out, _ = attn(x, position_embeddings=(cos, sin))
    assert out.shape == (1, 4, 68)
    assert torch.isfinite(out).all()

Classes

class DummyEmptyAttention

Expand source code

class DummyEmptyAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.foo = nn.Linear(64, 64)

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

class DummyGPT2Attention (dim=64)

Expand source code

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)

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

class DummyQwenAttention (dim=64, num_kv_heads=2, head_dim=16)

Expand source code

class DummyQwenAttention(nn.Module):
    def __init__(self, dim=64, num_kv_heads=2, head_dim=16):
        super().__init__()
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, num_kv_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_kv_heads * head_dim)
        self.o_proj = nn.Linear(dim, dim)

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

class DummyRobertaAttention (dim=64)

Expand source code

class DummyRobertaAttention(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.self = DummyRobertaSelfAttention(dim)
        self.output = nn.Module()
        self.output.dense = nn.Linear(dim, dim)

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

class DummyRobertaSelfAttention (dim=64)

Expand source code

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

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