Optimize computation orders

[pengwa]

Optimize computation orders

In Roberta/Electra, when ClassificationHead is used, there is slicing operation on features on sequence_length dimensions, then loss calculations only depend on this sliced data. This is a slicing at axis 1. Before slicing the shape is [batch, sequence_length, hidden], after slicing, it becomes [batch , hidden_stage]

We had opportunities to bring this slicing earlier as much as possible, by passing through simple elementwise ops (like Add/Div), or Layernorm/Softmax(if their reduce axis is after the slicing axis), or even MatMul's the left operand (if only it did not affect the last dims).

For operators like Reshape/Transpose, it is special since they have either data specified (after slicing we need update), or they have perm specified, which requires the input rank remain unchanged. So for those kinds of operators, we can remain the original rank, but just leave the sliced dim to be 1, after the compute completed, we do a Squeeze.

class RobertaClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""

    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, features, **kwargs):
        x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x

src\transformers\models\roberta\modeling_roberta.py
src\transformers\models\electra\modeling_electra.py

Benchmark

A simple benchmark shows Robeta training latency dropped from 208ms ~ 199ms. 4.5+% reduction.
More comprehensive tests are on the way.

Motivation and Context

[baijumeswani]

Looks good. Sorry for the delay in the review. Please update the branch to address failing pipelines.

[askhade]

Why are we deleting this transformer?

[pengwa]

Oh This is a renaming., now it is called ComputeOptimizer

[askhade]

All the code in these files and in the test files is wrapped in ENABLE_TRAINING so why add these files here? You can add them within the if (onnxruntime_ENABLE_TRAINING) condition

[pengwa]

Because the optimizer is applicable for inferencing. I intent to make it easier to enable it for inference later without changing file location and cmake macros.

[askhade]

what is the benefit of creating these const initializer list? Since any opportunity for reuse across the ops is purely coincidental.

[pengwa]

Defining a temporary list of opsets, on windows, later when we read it it is found some values are always zero. So I have to make is a constant to make it run correctly.

[askhade]

LGTM