⚡️ Speed up method IBertEmbeddings.forward by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for IBertEmbeddings.forward in src/transformers/models/ibert/modeling_ibert.py

⏱️ Runtime : 2.18 milliseconds → 2.06 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 6% speedup through three key micro-optimizations that reduce redundant tensor operations and attribute lookups:

1. Device Lookup Optimization
The original code repeatedly accessed .device attributes on tensors (e.g., input_ids.device, self.position_ids.device). The optimized version computes the device once at the beginning and reuses it, eliminating repeated attribute lookups which have small but measurable overhead in PyTorch.

2. Conditional Device Transfer
Instead of always calling .to(input_ids.device) on position_ids, the optimization adds a check if position_ids.device != device to avoid unnecessary device transfers when tensors are already on the correct device. This prevents redundant CUDA operations.

3. Tensor Operation Fusion in Position ID Creation
In create_position_ids_from_input_ids, the original code used .ne(padding_idx).int() which creates an intermediate boolean tensor. The optimized version uses input_ids != padding_idx followed by .int() on a separate line, allowing PyTorch to potentially optimize the operation chain better and reuse the mask_int variable.

Performance Impact
The line profiler shows the optimizations are most effective on error-handling test cases (7-10% improvements) where tensor operations dominate the runtime. For normal forward passes, the gains are more modest but consistent across different input sizes. These micro-optimizations are particularly valuable in transformer models where embeddings are called frequently during training and inference, making even small per-call improvements significant at scale.

The optimizations maintain identical functionality while reducing computational overhead through smarter tensor device management and operation fusion.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 43 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
import torch

from transformers.models.ibert.modeling_ibert import IBertEmbeddings


# Minimal config class for IBertEmbeddings
class DummyConfig:
    def __init__(
        self,
        vocab_size=100,
        hidden_size=32,
        pad_token_id=0,
        type_vocab_size=2,
        max_position_embeddings=128,
        layer_norm_eps=1e-12,
        hidden_dropout_prob=0.1,
        quant_mode=True,
        force_dequant=True,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.pad_token_id = pad_token_id
        self.type_vocab_size = type_vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.layer_norm_eps = layer_norm_eps
        self.hidden_dropout_prob = hidden_dropout_prob
        self.quant_mode = quant_mode
        self.force_dequant = force_dequant


# function to test (copied from above)
# ... [IBertEmbeddings, create_position_ids_from_input_ids as above] ...

# Basic Test Cases


def test_forward_with_inputs_embeds():
    # Test with inputs_embeds instead of input_ids
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 8
    inputs_embeds = torch.randn(batch_size, seq_len, config.hidden_size)
    embeddings, scaling = model.forward(inputs_embeds=inputs_embeds)


def test_forward_with_empty_input():
    # Test with empty input_ids (zero-length sequence)
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 0
    input_ids = torch.empty(batch_size, seq_len, dtype=torch.long)
    with pytest.raises(RuntimeError):
        # Should raise because sequence length is zero
        model.forward(input_ids=input_ids)  # 385μs -> 368μs (4.53% faster)


def test_forward_with_all_padding():
    # Test with input_ids consisting only of pad_token_id
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 5
    input_ids = torch.full((batch_size, seq_len), config.pad_token_id, dtype=torch.long)
    embeddings, scaling = model.forward(input_ids=input_ids)


def test_forward_with_max_position_ids():
    # Test with position_ids at maximum allowed value
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 1, config.max_position_embeddings
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    position_ids = torch.arange(config.max_position_embeddings).unsqueeze(0)
    embeddings, scaling = model.forward(input_ids=input_ids, position_ids=position_ids)


def test_forward_with_large_pad_token_id():
    # Test with pad_token_id set to a large value
    config = DummyConfig(pad_token_id=99)
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 10
    input_ids = torch.full((batch_size, seq_len), config.pad_token_id, dtype=torch.long)
    embeddings, scaling = model.forward(input_ids=input_ids)


def test_forward_with_inputs_embeds_and_position_ids():
    # Test with inputs_embeds and explicit position_ids
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 8
    inputs_embeds = torch.randn(batch_size, seq_len, config.hidden_size)
    position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, seq_len)
    embeddings, scaling = model.forward(inputs_embeds=inputs_embeds, position_ids=position_ids)


def test_forward_with_invalid_input_ids():
    # Test with input_ids containing values out of vocab range
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 8
    input_ids = torch.full((batch_size, seq_len), config.vocab_size + 1, dtype=torch.long)
    with pytest.raises(IndexError):
        # Should raise because input_ids are out of vocab range
        model.forward(input_ids=input_ids)  # 313μs -> 292μs (7.15% faster)


def test_forward_with_invalid_token_type_ids():
    # Test with token_type_ids out of type vocab range
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 8
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    token_type_ids = torch.full((batch_size, seq_len), config.type_vocab_size + 1, dtype=torch.long)
    with pytest.raises(IndexError):
        # Should raise because token_type_ids are out of range
        model.forward(input_ids=input_ids, token_type_ids=token_type_ids)  # 381μs -> 346μs (10.3% faster)


def test_forward_with_invalid_position_ids():
    # Test with position_ids out of range
    config = DummyConfig()
    model = IBertEmbeddings(config)
    batch_size, seq_len = 2, 8
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    position_ids = torch.full((batch_size, seq_len), config.max_position_embeddings + 1, dtype=torch.long)
    with pytest.raises(IndexError):
        # Should raise because position_ids are out of range
        model.forward(input_ids=input_ids, position_ids=position_ids)  # 732μs -> 687μs (6.55% faster)


def test_forward_with_mismatched_shapes():
    # Test with mismatched input_ids and token_type_ids shapes
    config = DummyConfig()
    model = IBertEmbeddings(config)
    input_ids = torch.randint(1, config.vocab_size, (2, 5), dtype=torch.long)
    token_type_ids = torch.randint(0, config.type_vocab_size, (3, 5), dtype=torch.long)
    with pytest.raises(RuntimeError):
        # Should raise due to shape mismatch
        model.forward(input_ids=input_ids, token_type_ids=token_type_ids)  # 370μs -> 362μs (2.09% faster)


# Large Scale Test Cases


def test_forward_large_batch_and_sequence():
    # Test with large batch and sequence length
    config = DummyConfig(vocab_size=1000, hidden_size=64, max_position_embeddings=512)
    model = IBertEmbeddings(config)
    batch_size, seq_len = 32, 256  # 32*256*64 = 524288 elements ~2MB
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    embeddings, scaling = model.forward(input_ids=input_ids)


def test_forward_maximum_tensor_size():
    # Test with maximum tensor size under 100MB
    config = DummyConfig(hidden_size=128)
    model = IBertEmbeddings(config)
    batch_size, seq_len = 64, 128  # 64*128*128 = 1048576 elements ~4MB
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    embeddings, scaling = model.forward(input_ids=input_ids)


def test_forward_large_type_vocab_size():
    # Test with large type_vocab_size
    config = DummyConfig(type_vocab_size=20)
    model = IBertEmbeddings(config)
    batch_size, seq_len = 3, 10
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), dtype=torch.long)
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_len), dtype=torch.long)
    embeddings, scaling = model.forward(input_ids=input_ids, token_type_ids=token_type_ids)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

# imports
import pytest
import torch

from transformers.models.ibert.modeling_ibert import IBertEmbeddings


# Mocks for required classes in IBertEmbeddings (minimal functional stubs)
class DummyConfig:
    def __init__(
        self,
        vocab_size=100,
        hidden_size=32,
        pad_token_id=0,
        type_vocab_size=2,
        max_position_embeddings=128,
        layer_norm_eps=1e-12,
        hidden_dropout_prob=0.1,
        quant_mode=True,
        force_dequant=False,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.pad_token_id = pad_token_id
        self.type_vocab_size = type_vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.layer_norm_eps = layer_norm_eps
        self.hidden_dropout_prob = hidden_dropout_prob
        self.quant_mode = quant_mode
        self.force_dequant = force_dequant


# -------------------- UNIT TESTS --------------------

# Basic Test Cases


def test_forward_eval_no_dropout():
    # Test that dropout is disabled in eval mode
    config = DummyConfig(
        vocab_size=10, hidden_size=8, pad_token_id=0, max_position_embeddings=20, hidden_dropout_prob=0.5
    )
    model = IBertEmbeddings(config)
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
    model.eval()
    outputs1, _ = model(input_ids=input_ids)
    outputs2, _ = model(input_ids=input_ids)


# Error Handling Test Cases


def test_forward_invalid_token_type_ids_shape():
    # Test with token_type_ids shape mismatch
    config = DummyConfig(vocab_size=10, hidden_size=8, pad_token_id=0, type_vocab_size=2, max_position_embeddings=20)
    model = IBertEmbeddings(config)
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
    token_type_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)  # wrong shape
    with pytest.raises(RuntimeError):
        model(input_ids=input_ids, token_type_ids=token_type_ids)


def test_forward_invalid_position_ids_shape():
    # Test with position_ids shape mismatch
    config = DummyConfig(vocab_size=10, hidden_size=8, pad_token_id=0, max_position_embeddings=20)
    model = IBertEmbeddings(config)
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
    position_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)  # wrong shape
    with pytest.raises(RuntimeError):
        model(input_ids=input_ids, position_ids=position_ids)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-IBertEmbeddings.forward-mia3od7a and push.