Request for adding supports of Ops `XlaDynamicUpdateSlice` and `XlaDynamicSlice`

[csAugust]

Hello TF.js team. Thanks for your effort on helping deal with conversion among different ML model formats, but I'll be more appreciated if someone could add the support for Ops XlaDynamicUpdateSlice and XlaDynamicSlice accounting for my failure when I tried to convert some models of savedModel format.

System information

• Python 3.9.17
• TensorFlow version: 2.13.0
• tensorflowjs library version: 4.11.0

Describe the feature and the current behavior/state.
I want to run some language models for text generation tasks in browser so I first searched three models(flan-t5-base, gpt2 and dialoGPT) from HuggingFace and saved them in savedModel format. But when I tried to convert any of these models of savedModel format, I got an error:

ValueError: Unsupported Ops in the model before optimization
XlaDynamicUpdateSlice

or

ValueError: Unsupported Ops in the model before optimization
XlaDynamicSlice, XlaDynamicUpdateSlice

So I think whether someone can help with adding these features to support these two Ops such that I can successfully convert these models.

Will this change the current api? How?
I think it won't as it could should as an inner function.

Who will benefit with this feature?
As I have tried several language models for text generation but all failed, it will be much beneficial for all who needs this kind of models running in tfjs once this feature be added.

Any Other info.
I ran the converter script by:

tensorflowjs_converter \
--input_format=tf_saved_model \
./original_model_path \
./target_model_path

Meanwhile, if anyone knows there are any possible language model candidates for text generation of high quality that can be run (or can be run after converted) in tfjs, please tells me more and I'll appreciate your kindness.

[gaikwadrahul8]

Hi, @csAugust

Thank you for bringing this issue to our attention and if possible could you please help us with Google colab notebook where you're getting ValueError: Unsupported Ops XlaDynamicSlice, XlaDynamicUpdateSlice to replicate the same behaviour from our end

Most probably those are Unsupported Ops while converting and you can refer official documentation for supported Ops in tfjs-converter. Thank you!

[csAugust]

Hi, @csAugust

Thank you for bringing this issue to our attention and if possible could you please help us with Google colab notebook where you're getting ValueError: Unsupported Ops XlaDynamicSlice, XlaDynamicUpdateSlice to replicate the same behaviour from our end

Most probably those are Unsupported Ops while converting and you can refer official documentation for supported Ops in tfjs-converter. Thank you!

Of course and here is the link to a colab script containing how I got this error: https://colab.research.google.com/drive/1U0oGXrhgO3Cax0on1_Rj9kXdz7RAiSVG?usp=sharing
however, unfortunately it seems that it can't be run in colab due to memory allocation problems.

I don't know if I have made some mistakes when transfer huggingface models to tfjs in any sections. I checked the docs you mentioned and I'm confused that it seems like XlaDynamicSlice, XlaDynamicUpdateSlice aren't tensorflow ops. So I can't really tell what causes this bug, and would you please help me out? Thanks a lot!

[gaikwadrahul8]

Hi, @csAugust

I tried to replicate the same issue from my end and I'm also getting the same behaviour in Google colab, here is gist-file for reference and As far I know and by looking at the error messages above these are Unsupported Ops XlaDynamicSlice, XlaDynamicUpdateSlice while converting huggingface model into TFJs so this issue will be considered as feature request and at the moment I am unable to provide an exact ETA for when these Ops will be implemented by our development team, as it depends on priority and bandwidth.

In the meantime, if anyone would like to contribute to the implementation of the Unsupported Ops XlaDynamicSlice and XlaDynamicUpdateSlice, please feel free to do so. Please refer to the following links for more information Ref-1, Ref-2.

Thank you!

[mattsoulanille]

We've run into the same issue with these ops when we tried converting similar LLMs. They're XLA ops instead of TensorFlow ops, so we don't have any precedent for implementing them, although we may need to in the future. If you'd like, you can add a custom TFJS op for each of them. In fact, we have a not-perfectly-efficient but still potentially useful implementation of XlaDynamicUpdateSlice in tfjs-layers that you could potentially register as a custom op.

For other LLMs to try, I've had some success converting Facebook's OPT models, which don't seem to require these ops.

[csAugust]

We've run into the same issue with these ops when we tried converting similar LLMs. They're XLA ops instead of TensorFlow ops, so we don't have any precedent for implementing them, although we may need to in the future. If you'd like, you can add a custom TFJS op for each of them. In fact, we have a not-perfectly-efficient but still potentially useful implementation of XlaDynamicUpdateSlice in tfjs-layers that you could potentially register as a custom op.

For other LLMs to try, I've had some success converting Facebook's OPT models, which don't seem to require these ops.

Thanks for your attention and help! Before I try to implement it as a custom op, I tried converting OPT models you mentioned. But to my confuse I have met another error when saving the model to SavedModel format. It likes:

TypeError: in user code:

    File "/home/zhiyangchen/code/Tensorflow/opt-350m/opt.py", line 11, in generate  *
        outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask)
    File "/home/zhiyangchen/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 911, in generate  *
        **model_kwargs,
    File "/home/zhiyangchen/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 1703, in greedy_search_body_fn  *
        model_kwargs = self._update_model_kwargs_for_xla_generation(
    File "/home/zhiyangchen/miniconda3/envs/py39/lib/python3.9/site-packages/transformers/generation/tf_utils.py", line 1348, in _initialize_past  *
        padding_values = tf.constant([[0, 0], [0, 0], [0, num_padding_values], [0, 0]], dtype=tf.int32)

    TypeError: Expected int32, but got Tensor("sub_4:0", shape=(), dtype=int32) of type 'Tensor'.

I wonder if my code of defining signatures is wrong or anything else:

@tf.function(input_signature=[tf.TensorSpec(shape=[1, None], dtype=tf.int32, name="input_ids"),
                              tf.TensorSpec(shape=[1, None], dtype=tf.int32, name="attention_mask")])
def generate(input_ids, attention_mask):
    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask)
    return {"generated": outputs}

tf.saved_model.save(model, "opt-350m_save", signatures={"serving_default": generate})

Sorry to bother but could you please tell me more about saving the huggingface models? I've searched some docs for help but I still can't properly deal with it :( . Thank you!