Module neuralops.tests.test_models_and_datasets_matrix

Comprehensive Matrix Test Suite: Small Models (BERT, RoBERTa, DistilBERT, GPT-2) across all Datasets (GLUE tasks, WikiText-2, C4, Multi-Choice) and SOTA Optimizers/Compressors.

Verifies 100% end-to-end compatibility with zero conflicts: 1. Model Architectures: BERT, RoBERTa, DistilBERT, GPT-2 causal LM. 2. Datasets: GLUE (sst2, mrpc, qnli, stsb, cola, rte, mnli), WikiText-2, C4, Multi-Choice. 3. Compression: Dynamic Entropy SVD, Res-SVD, AASVD, SliceGPT, Kronecker. 4. Optimization: Riemannian Stiefel Optimizer, Factorized Gradient Clipping, Mixed Precision AMP.

Classes

class CompactBertForClassification (vocab_size: int = 1000,
hidden_size: int = 64,
num_layers: int = 2,
num_heads: int = 4,
num_labels: int = 2)

Expand source code

class CompactBertForClassification(nn.Module):
    """Full BERT / RoBERTa classifier model for sequence classification benchmarks."""

    def __init__(
        self,
        vocab_size: int = 1000,
        hidden_size: int = 64,
        num_layers: int = 2,
        num_heads: int = 4,
        num_labels: int = 2,
    ) -> None:
        super().__init__()
        self.embeddings = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList(
            [CompactBertLayer(hidden_size, num_heads) for _ in range(num_layers)]
        )
        self.pooler = nn.Linear(hidden_size, hidden_size)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> Dict[str, torch.Tensor]:
        x = self.embeddings(input_ids)
        for layer in self.layers:
            x = layer(x, attention_mask=attention_mask)
        pooled = torch.tanh(self.pooler(x[:, 0]))  # [CLS] token representation
        logits = self.classifier(pooled)

        loss = None
        if labels is not None:
            if logits.shape[-1] == 1:
                loss = F.mse_loss(logits.squeeze(-1), labels.float())
            else:
                loss = F.cross_entropy(logits, labels.long())

        return {"logits": logits, "loss": loss if loss is not None else torch.tensor(0.0)}

Full BERT / RoBERTa classifier model for sequence classification benchmarks.

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
**kwargs: Any) ‑> Dict[str, torch.Tensor]

Expand source code

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    **kwargs: Any,
) -> Dict[str, torch.Tensor]:
    x = self.embeddings(input_ids)
    for layer in self.layers:
        x = layer(x, attention_mask=attention_mask)
    pooled = torch.tanh(self.pooler(x[:, 0]))  # [CLS] token representation
    logits = self.classifier(pooled)

    loss = None
    if labels is not None:
        if logits.shape[-1] == 1:
            loss = F.mse_loss(logits.squeeze(-1), labels.float())
        else:
            loss = F.cross_entropy(logits, labels.long())

    return {"logits": logits, "loss": loss if loss is not None else torch.tensor(0.0)}

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class CompactBertLayer (hidden_size: int = 64, num_heads: int = 4)

Expand source code

class CompactBertLayer(nn.Module):
    """Complete BERT / RoBERTa transformer layer."""

    def __init__(self, hidden_size: int = 64, num_heads: int = 4) -> None:
        super().__init__()
        self.attention = CompactBertSelfAttention(hidden_size, num_heads)
        self.norm1 = nn.LayerNorm(hidden_size)
        self.intermediate = nn.Linear(hidden_size, hidden_size * 2)
        self.output = nn.Linear(hidden_size * 2, hidden_size)
        self.norm2 = nn.LayerNorm(hidden_size)

    def forward(
        self, x: torch.Tensor, attention_mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        attn_out = self.attention(self.norm1(x), attention_mask)
        x = x + attn_out
        mlp_out = self.output(F.gelu(self.intermediate(self.norm2(x))))
        return x + mlp_out

Complete BERT / RoBERTa transformer layer.

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def forward(
    self, x: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
    attn_out = self.attention(self.norm1(x), attention_mask)
    x = x + attn_out
    mlp_out = self.output(F.gelu(self.intermediate(self.norm2(x))))
    return x + mlp_out

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class CompactBertSelfAttention (hidden_size: int = 64, num_heads: int = 4)

Expand source code

class CompactBertSelfAttention(nn.Module):
    """BERT / RoBERTa self-attention layer with separate Q, K, V projections."""

    def __init__(self, hidden_size: int = 64, num_heads: int = 4) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.query = nn.Linear(hidden_size, hidden_size)
        self.key = nn.Linear(hidden_size, hidden_size)
        self.value = nn.Linear(hidden_size, hidden_size)
        self.dense = nn.Linear(hidden_size, hidden_size)

    def forward(
        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        B, N, D = hidden_states.shape
        q = self.query(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = self.key(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.value(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                scores = scores + (1.0 - attention_mask[:, None, None, :]) * -1e4
        probs = torch.softmax(scores, dim=-1)
        context = torch.matmul(probs, v).permute(0, 2, 1, 3).contiguous().view(B, N, D)
        return self.dense(context)

BERT / RoBERTa self-attention layer with separate Q, K, V projections.

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) ‑> torch.Tensor

Expand source code

def forward(
    self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
    B, N, D = hidden_states.shape
    q = self.query(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
    k = self.key(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
    v = self.value(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

    scores = torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            scores = scores + (1.0 - attention_mask[:, None, None, :]) * -1e4
    probs = torch.softmax(scores, dim=-1)
    context = torch.matmul(probs, v).permute(0, 2, 1, 3).contiguous().view(B, N, D)
    return self.dense(context)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class CompactGPT2ForCausalLM (vocab_size: int = 1000,
hidden_size: int = 64,
num_layers: int = 2,
num_heads: int = 4)

Expand source code

class CompactGPT2ForCausalLM(nn.Module):
    """Full GPT-2 causal language model with fused QKV and decay-gated attention options."""

    def __init__(
        self,
        vocab_size: int = 1000,
        hidden_size: int = 64,
        num_layers: int = 2,
        num_heads: int = 4,
    ) -> None:
        super().__init__()
        self.wte = nn.Embedding(vocab_size, hidden_size)
        config = SimpleNamespace(
            hidden_size=hidden_size,
            num_attention_heads=num_heads,
            dropout=0.0,
            linear_attention_eps=1e-6,
            use_decay_gating=True,
            use_fused_kernel=True,
        )
        self.h = nn.ModuleList(
            [
                SVDLinearAttention(config, rank=16, mode="decay", use_decay_gating=True)
                for _ in range(num_layers)
            ]
        )
        self.ln_f = nn.LayerNorm(hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> Dict[str, torch.Tensor]:
        x = self.wte(input_ids)
        for block in self.h:
            attn_out, _ = block(x, attention_mask=attention_mask)
            x = x + attn_out
        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
            )

        return {"logits": logits, "loss": loss if loss is not None else torch.tensor(0.0)}

Full GPT-2 causal language model with fused QKV and decay-gated attention options.

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

Ancestors

• torch.nn.modules.module.Module

Methods

def forward(self,
input_ids: torch.Tensor,
attention_mask: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
**kwargs: Any) ‑> Dict[str, torch.Tensor]

Expand source code

def forward(
    self,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    **kwargs: Any,
) -> Dict[str, torch.Tensor]:
    x = self.wte(input_ids)
    for block in self.h:
        attn_out, _ = block(x, attention_mask=attention_mask)
        x = x + attn_out
    x = self.ln_f(x)
    logits = self.lm_head(x)

    loss = None
    if labels is not None:
        shift_logits = logits[:, :-1, :].contiguous()
        shift_labels = labels[:, 1:].contiguous()
        loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
        )

    return {"logits": logits, "loss": loss if loss is not None else torch.tensor(0.0)}

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class TestModelsAndDatasetsMatrix

Expand source code

class TestModelsAndDatasetsMatrix:
    """Comprehensive test matrix across models, GLUE datasets, language modeling, and compressors."""

    def test_bert_classification_on_glue_sst2(self) -> None:
        """Verify BERT classification model trains on GLUE SST-2 dataset."""
        torch.manual_seed(42)
        model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_labels=2)
        optimizer = RiemannianStiefelOptimizer(model.parameters(), lr=1e-3, base_optimizer="adamw")

        # Load synthetic SST-2 dataset
        dataset = SyntheticClassificationDataset(num_samples=16, seq_len=16, vocab_size=500, num_classes=2)
        loader = DataLoader(dataset, batch_size=4)

        for batch in loader:
            optimizer.zero_grad()
            out = model(**batch)
            loss = out["loss"]
            loss.backward()
            clip_grad_norm_factorized(model.parameters(), max_norm=1.0)
            optimizer.step()

        assert torch.isfinite(loss).all()

    def test_roberta_regression_on_glue_stsb(self) -> None:
        """Verify RoBERTa regression model on continuous GLUE STS-B task."""
        torch.manual_seed(42)
        model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_labels=1)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

        input_ids = torch.randint(1, 500, (4, 16))
        attention_mask = torch.ones((4, 16))
        labels = torch.tensor([1.2, 3.5, 4.8, 0.5])

        out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        assert out["logits"].shape == (4, 1)
        assert torch.isfinite(out["loss"]).all()

        out["loss"].backward()
        optimizer.step()

    def test_distilbert_with_dynamic_entropy_svd(self) -> None:
        """Verify DistilBERT compresses via Dynamic Entropy SVD and executes inference."""
        torch.manual_seed(42)
        model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_layers=2)
        dataset = SyntheticClassificationDataset(num_samples=16, seq_len=16, vocab_size=500, num_classes=2)
        loader = DataLoader(dataset, batch_size=4)

        compressor = DynamicEntropySVDCompressor(
            model,
            target_rank_budget=64,
            min_rank=8,
            max_rank=32,
            calibration_loader=loader,
            use_fused_kernel=True,
        )
        compressed_model, stats = compressor.compress()
        assert stats["compressed_parameters"] <= stats["original_parameters"]

        # Run inference on compressed DistilBERT
        sample = dataset[0]
        input_ids = sample["input_ids"].unsqueeze(0)
        out = compressed_model(input_ids)
        assert out["logits"].shape == (1, 2)
        assert torch.isfinite(out["logits"]).all()

    def test_gpt2_causal_lm_on_wikitext2_and_c4(self) -> None:
        """Verify GPT-2 causal language model with decay linear attention trains on WikiText-2 and C4."""
        torch.manual_seed(42)
        model = CompactGPT2ForCausalLM(vocab_size=500, hidden_size=64, num_layers=2, num_heads=4)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

        # Load WikiText-2 synthetic benchmark
        w2_samples = load_benchmark_dataset(
            "wikitext2", use_synthetic=True, max_samples=8, return_dataloader=False
        )
        assert len(w2_samples) == 8

        # Train on sequence batch
        tokens = generate_zipfian_tokens(num_samples=4, seq_len=32, vocab_size=500)
        optimizer.zero_grad()
        out = model(input_ids=tokens, labels=tokens)
        assert out["logits"].shape == (4, 32, 500)
        assert out["loss"] > 0.0

        out["loss"].backward()
        clip_grad_norm_factorized(model.parameters(), max_norm=1.0)
        optimizer.step()

    def test_bert_res_svd_spectral_residual_lora(self) -> None:
        """Verify Spectral Residual LoRA on BERT freezes base and updates adapter parameters."""
        torch.manual_seed(42)
        model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_layers=2)
        compressor = ResSVDCompressor(
            model,
            base_rank=12,
            residual_rank=4,
            lora_alpha=8.0,
            use_fused_kernel=True,
        )
        res_model, stats = compressor.compress()
        assert len(stats["replaced_layers"]) > 0

        # Fine-tune adapter parameters
        optimizer = torch.optim.AdamW(
            [p for p in res_model.parameters() if p.requires_grad], lr=1e-3
        )
        input_ids = torch.randint(1, 500, (4, 16))
        labels = torch.tensor([0, 1, 0, 1])

        out = res_model(input_ids=input_ids, labels=labels)
        out["loss"].backward()
        optimizer.step()
        assert torch.isfinite(out["loss"]).all()

    def test_multi_choice_dataset_pipeline(self) -> None:
        """Verify MultiChoiceDataset loading and tensor shaping for ARC/MMLU/HellaSwag tasks."""
        import pandas as pd
        torch.manual_seed(42)

        data = {
            "question": [
                "Which element has atomic number 1?",
                "What is the speed of light in vacuum?",
            ],
            "choices": [
                ["Hydrogen", "Helium", "Lithium", "Beryllium"],
                ["300 km/s", "300,000 km/s", "150,000 km/s", "Infinity"],
            ],
            "label": [0, 1],
        }
        df = pd.DataFrame(data)

        def mock_tokenizer(text, truncation=True, padding=True, max_length=16, return_tensors="pt"):
            return {
                "input_ids": torch.randint(1, 500, (1, max_length)),
                "attention_mask": torch.ones((1, max_length)),
            }

        dataset = MultiChoiceDataset(df=df, tokenizer=mock_tokenizer, max_length=16)
        loader = DataLoader(dataset, batch_size=2)
        batch = next(iter(loader))

        assert "input_ids" in batch
        assert "labels" in batch
        assert batch["input_ids"].shape == (2, 4, 16)
        assert batch["labels"].shape == (2,)
        assert (batch["labels"] == torch.tensor([0, 1])).all()

    def test_fused_qkv_slicing_on_all_attention_heads(self) -> None:
        """Verify slice_fused_qkv and split_and_set_qkv_weights for standard MHA, GQA, and MQA."""
        torch.manual_seed(42)
        hidden_size = 64
        num_heads = 4
        head_dim = 16

        # Standard MHA: [3 * hidden_size, hidden_size]
        fused_w_mha = torch.randn(3 * hidden_size, hidden_size)
        fused_b_mha = torch.randn(3 * hidden_size)
        qkv_mha = slice_fused_qkv(fused_w_mha, fused_b_mha, num_heads=num_heads, head_dim=head_dim)
        assert qkv_mha["q"][0].shape == (hidden_size, hidden_size)
        assert qkv_mha["k"][0].shape == (hidden_size, hidden_size)
        assert qkv_mha["v"][0].shape == (hidden_size, hidden_size)

        # GQA: 4 query heads, 2 key/value heads
        fused_w_gqa = torch.randn(4 * head_dim + 2 * head_dim + 2 * head_dim, hidden_size)
        qkv_gqa = slice_fused_qkv(
            fused_w_gqa, None, num_heads=4, num_kv_heads=2, head_dim=head_dim
        )
        assert qkv_gqa["q"][0].shape == (64, hidden_size)
        assert qkv_gqa["k"][0].shape == (32, hidden_size)
        assert qkv_gqa["v"][0].shape == (32, hidden_size)

Comprehensive test matrix across models, GLUE datasets, language modeling, and compressors.

Methods

def test_bert_classification_on_glue_sst2(self) ‑> None

Expand source code

def test_bert_classification_on_glue_sst2(self) -> None:
    """Verify BERT classification model trains on GLUE SST-2 dataset."""
    torch.manual_seed(42)
    model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_labels=2)
    optimizer = RiemannianStiefelOptimizer(model.parameters(), lr=1e-3, base_optimizer="adamw")

    # Load synthetic SST-2 dataset
    dataset = SyntheticClassificationDataset(num_samples=16, seq_len=16, vocab_size=500, num_classes=2)
    loader = DataLoader(dataset, batch_size=4)

    for batch in loader:
        optimizer.zero_grad()
        out = model(**batch)
        loss = out["loss"]
        loss.backward()
        clip_grad_norm_factorized(model.parameters(), max_norm=1.0)
        optimizer.step()

    assert torch.isfinite(loss).all()

Verify BERT classification model trains on GLUE SST-2 dataset.

def test_bert_res_svd_spectral_residual_lora(self) ‑> None

Expand source code

def test_bert_res_svd_spectral_residual_lora(self) -> None:
    """Verify Spectral Residual LoRA on BERT freezes base and updates adapter parameters."""
    torch.manual_seed(42)
    model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_layers=2)
    compressor = ResSVDCompressor(
        model,
        base_rank=12,
        residual_rank=4,
        lora_alpha=8.0,
        use_fused_kernel=True,
    )
    res_model, stats = compressor.compress()
    assert len(stats["replaced_layers"]) > 0

    # Fine-tune adapter parameters
    optimizer = torch.optim.AdamW(
        [p for p in res_model.parameters() if p.requires_grad], lr=1e-3
    )
    input_ids = torch.randint(1, 500, (4, 16))
    labels = torch.tensor([0, 1, 0, 1])

    out = res_model(input_ids=input_ids, labels=labels)
    out["loss"].backward()
    optimizer.step()
    assert torch.isfinite(out["loss"]).all()

Verify Spectral Residual LoRA on BERT freezes base and updates adapter parameters.

def test_distilbert_with_dynamic_entropy_svd(self) ‑> None

Expand source code

def test_distilbert_with_dynamic_entropy_svd(self) -> None:
    """Verify DistilBERT compresses via Dynamic Entropy SVD and executes inference."""
    torch.manual_seed(42)
    model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_layers=2)
    dataset = SyntheticClassificationDataset(num_samples=16, seq_len=16, vocab_size=500, num_classes=2)
    loader = DataLoader(dataset, batch_size=4)

    compressor = DynamicEntropySVDCompressor(
        model,
        target_rank_budget=64,
        min_rank=8,
        max_rank=32,
        calibration_loader=loader,
        use_fused_kernel=True,
    )
    compressed_model, stats = compressor.compress()
    assert stats["compressed_parameters"] <= stats["original_parameters"]

    # Run inference on compressed DistilBERT
    sample = dataset[0]
    input_ids = sample["input_ids"].unsqueeze(0)
    out = compressed_model(input_ids)
    assert out["logits"].shape == (1, 2)
    assert torch.isfinite(out["logits"]).all()

Verify DistilBERT compresses via Dynamic Entropy SVD and executes inference.

def test_fused_qkv_slicing_on_all_attention_heads(self) ‑> None

Expand source code

def test_fused_qkv_slicing_on_all_attention_heads(self) -> None:
    """Verify slice_fused_qkv and split_and_set_qkv_weights for standard MHA, GQA, and MQA."""
    torch.manual_seed(42)
    hidden_size = 64
    num_heads = 4
    head_dim = 16

    # Standard MHA: [3 * hidden_size, hidden_size]
    fused_w_mha = torch.randn(3 * hidden_size, hidden_size)
    fused_b_mha = torch.randn(3 * hidden_size)
    qkv_mha = slice_fused_qkv(fused_w_mha, fused_b_mha, num_heads=num_heads, head_dim=head_dim)
    assert qkv_mha["q"][0].shape == (hidden_size, hidden_size)
    assert qkv_mha["k"][0].shape == (hidden_size, hidden_size)
    assert qkv_mha["v"][0].shape == (hidden_size, hidden_size)

    # GQA: 4 query heads, 2 key/value heads
    fused_w_gqa = torch.randn(4 * head_dim + 2 * head_dim + 2 * head_dim, hidden_size)
    qkv_gqa = slice_fused_qkv(
        fused_w_gqa, None, num_heads=4, num_kv_heads=2, head_dim=head_dim
    )
    assert qkv_gqa["q"][0].shape == (64, hidden_size)
    assert qkv_gqa["k"][0].shape == (32, hidden_size)
    assert qkv_gqa["v"][0].shape == (32, hidden_size)

Verify slice_fused_qkv and split_and_set_qkv_weights for standard MHA, GQA, and MQA.

def test_gpt2_causal_lm_on_wikitext2_and_c4(self) ‑> None

Expand source code

def test_gpt2_causal_lm_on_wikitext2_and_c4(self) -> None:
    """Verify GPT-2 causal language model with decay linear attention trains on WikiText-2 and C4."""
    torch.manual_seed(42)
    model = CompactGPT2ForCausalLM(vocab_size=500, hidden_size=64, num_layers=2, num_heads=4)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    # Load WikiText-2 synthetic benchmark
    w2_samples = load_benchmark_dataset(
        "wikitext2", use_synthetic=True, max_samples=8, return_dataloader=False
    )
    assert len(w2_samples) == 8

    # Train on sequence batch
    tokens = generate_zipfian_tokens(num_samples=4, seq_len=32, vocab_size=500)
    optimizer.zero_grad()
    out = model(input_ids=tokens, labels=tokens)
    assert out["logits"].shape == (4, 32, 500)
    assert out["loss"] > 0.0

    out["loss"].backward()
    clip_grad_norm_factorized(model.parameters(), max_norm=1.0)
    optimizer.step()

Verify GPT-2 causal language model with decay linear attention trains on WikiText-2 and C4.

def test_multi_choice_dataset_pipeline(self) ‑> None

Expand source code

def test_multi_choice_dataset_pipeline(self) -> None:
    """Verify MultiChoiceDataset loading and tensor shaping for ARC/MMLU/HellaSwag tasks."""
    import pandas as pd
    torch.manual_seed(42)

    data = {
        "question": [
            "Which element has atomic number 1?",
            "What is the speed of light in vacuum?",
        ],
        "choices": [
            ["Hydrogen", "Helium", "Lithium", "Beryllium"],
            ["300 km/s", "300,000 km/s", "150,000 km/s", "Infinity"],
        ],
        "label": [0, 1],
    }
    df = pd.DataFrame(data)

    def mock_tokenizer(text, truncation=True, padding=True, max_length=16, return_tensors="pt"):
        return {
            "input_ids": torch.randint(1, 500, (1, max_length)),
            "attention_mask": torch.ones((1, max_length)),
        }

    dataset = MultiChoiceDataset(df=df, tokenizer=mock_tokenizer, max_length=16)
    loader = DataLoader(dataset, batch_size=2)
    batch = next(iter(loader))

    assert "input_ids" in batch
    assert "labels" in batch
    assert batch["input_ids"].shape == (2, 4, 16)
    assert batch["labels"].shape == (2,)
    assert (batch["labels"] == torch.tensor([0, 1])).all()

Verify MultiChoiceDataset loading and tensor shaping for ARC/MMLU/HellaSwag tasks.

def test_roberta_regression_on_glue_stsb(self) ‑> None

Expand source code

def test_roberta_regression_on_glue_stsb(self) -> None:
    """Verify RoBERTa regression model on continuous GLUE STS-B task."""
    torch.manual_seed(42)
    model = CompactBertForClassification(vocab_size=500, hidden_size=64, num_labels=1)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

    input_ids = torch.randint(1, 500, (4, 16))
    attention_mask = torch.ones((4, 16))
    labels = torch.tensor([1.2, 3.5, 4.8, 0.5])

    out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
    assert out["logits"].shape == (4, 1)
    assert torch.isfinite(out["loss"]).all()

    out["loss"].backward()
    optimizer.step()

Verify RoBERTa regression model on continuous GLUE STS-B task.