Unnecessary GATHER_NDs that are then unsupported on LiteRT GPU Delegate

[BmanClark]

Description of the bug:

Bug is harsh - it works on CPU - but even there it's less efficient than it could be.
Example model:

model_path = "roblox/voice-safety-classifier-v2"
model = WavLMForSequenceClassification.from_pretrained(
    model_path, num_labels=len(labels_name_list)
)
dummy_input = (torch.randn(1, 80000),) 
edge_model = litert_torch.convert(model, dummy_input)
edge_model.export("./out.tflite")

Creates 2 GATHER_NDs early in the model that are in fact just doing RESHAPE from [512] to [1,512]. That section could all be simplified further as well (lots of unnecessary reshapes are generated), but biggest issue is the GATHER_NDs, which prevent GPU use. (There are other things in model that prevent GPU use, but they have other fixes)

Actual vs expected behavior:

Actual: GATHER_NDs unnecessarily created.
Expected: Where possible GPU supported ops are preferred - if a RESHAPE, TRANSPOSE or STRIDED_SLICE would do, prefer those.

Any other information you'd like to share?

GATHER_ND is not expected to be supported on GPU in near future (google-ai-edge/LiteRT#5197)

[sourcelite]

Hi @BmanClark, welcome back. you're correct - it isn't the model is not working, but it is that the converter is being too literal with the indexed access resulting in GATHER_ND ops that is killing GPU Delegation.
(litert-torch is like a bridge or a translator that takes PyTorch ATen code and tries to map it to LiteRT)

I did the root cause analysis - why does GATHER_ND appears - the slicing/indexing operations in some models, maybe like WavLM or Wav2Vec2, are interpreted by the litert-torch converter as multidimensional gather. So when you see the GATHER_ND performing [512] to [1,512] transition it is a RESHAPE. Again, as we discussed in #5197 - (google-ai-edge/LiteRT#5197), the gather_nd was GPU Killer becuase it forced random memory access pattern that break GPU texture cache and also thread synchrony.

I will shortly come back on this with info.

[sourcelite]

I predict it to be a good optimiser issue, we need to find why Gather_ND is not being folded into RESHAPE during the final export pass.

Thus, I need to see the snippet of specific layer in the WavLM model where this occurs - is it in feature extractor or in classification head, it will help me check on specific graph transformation rule to catch the pattern observed.
Looking for above information from you. Thanks!

[BmanClark]

Hello from the old issue! :)

The issue is in the GroupNorm, it's very early on in the model (ends up operators 23 & 25 in my Netron view of the converted model)

I'm loading the model through HuggingFace transformers, as above rather than from the original code, and then having to monkey patch fixes on the loaded model, but I think the code for WavLM is at https://github.com/microsoft/unilm/tree/master/wavlm and I can try and find the GroupNorm there if needed. But hopefully, the input/output matrices might be enough to identify?

The inputs to the first GATHER_ND are transformers.models.wavlm.modeling_wavlm.WavLMForSequenceClassification/transformers.models.wavlm.modeling_wavlm.WavLMModel_wavlm/transformers.models.wavlm.modeling_wavlm.WavLMFeatureEncoder_feature_extractor/transformers.models.wavlm.modeling_wavlm.WavLMGroupNormConvLayer_0/torch.nn.modules.normalization.GroupNorm_layer_norm;22 and transformers.models.wavlm.modeling_wavlm.WavLMForSequenceClassification/transformers.models.wavlm.modeling_wavlm.WavLMModel_wavlm/transformers.models.wavlm.modeling_wavlm.WavLMFeatureEncoder_feature_extractor/transformers.models.wavlm.modeling_wavlm.WavLMGroupNormConvLayer_0/torch.nn.modules.normalization.GroupNorm_layer_norm;2 and the output is transformers.models.wavlm.modeling_wavlm.WavLMForSequenceClassification/transformers.models.wavlm.modeling_wavlm.WavLMModel_wavlm/transformers.models.wavlm.modeling_wavlm.WavLMFeatureEncoder_feature_extractor/transformers.models.wavlm.modeling_wavlm.WavLMGroupNormConvLayer_0/torch.nn.modules.normalization.GroupNorm_layer_norm;23
The other GATHER_ND is parallel to it in the graph with very similar names, just slightly different final numbers and probably comes from the same code.

Looking at that original WavLM code it looks like the GroupNorms are just created from torch.nn.functional.group_norm

(I've actually got around this issue by monkey patching a much cleaner Group Norm over it, which didn't got to Gathers and indeed reduced the number of ops for that section massively while keeping accuracy, so I guess it's definitely somewhat "the code's fault", but if that code is a torch library then...)