Feature Request: Add support for embedding layers

[dopc]

Hey,
I already created an issue on Huggingface.
One of the issues can be closed if you need to.

I want to use concerete-ml for the Transformer model, such as BERT.
Do you have any resources to look at or advice you could give for this?

I already tried the distilbert ner and conll2003 by duplicating and modifying this model, however, I have not succeeded yet.

Thanks,
Best.

[jfrery]

Hi @dopc,

Could you share the error you are getting?

Some parts from the transformer architecture might be missing (e.g. softmax). They will be supported soon.

For now you can check the use case available on LLMs. It's using GPT2: https://github.com/zama-ai/concrete-ml/tree/main/use_case_examples/llm.

[dopc]

Hey @jfrery, thanks for your quick reply.

Good to see GPT2 model works! I will look at it.

Here is my error trace:

  File "/opt/conda/lib/python3.10/site-packages/concrete/ml/torch/compile.py", line 303, in compile_torch_model
    return _compile_torch_or_onnx_model(
  File "/opt/conda/lib/python3.10/site-packages/concrete/ml/torch/compile.py", line 195, in _compile_torch_or_onnx_model
    quantized_module = build_quantized_module(
  File "/opt/conda/lib/python3.10/site-packages/concrete/ml/torch/compile.py", line 117, in build_quantized_module
    numpy_model = NumpyModule(model, dummy_input_for_tracing)
  File "/opt/conda/lib/python3.10/site-packages/concrete/ml/torch/numpy_module.py", line 45, in __init__
    self.numpy_forward, self._onnx_model = get_equivalent_numpy_forward_from_torch(
  File "/opt/conda/lib/python3.10/site-packages/concrete/ml/onnx/convert.py", line 145, in get_equivalent_numpy_forward_from_torch
    torch.onnx.export(
  File "/opt/conda/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
    _export(
  File "/opt/conda/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/opt/conda/lib/python3.10/site-packages/torch/onnx/utils.py", line 1111, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/opt/conda/lib/python3.10/site-packages/torch/onnx/utils.py", line 987, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/opt/conda/lib/python3.10/site-packages/torch/onnx/utils.py", line 891, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "/opt/conda/lib/python3.10/site-packages/torch/jit/_trace.py", line 1184, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/jit/_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/opt/conda/lib/python3.10/site-packages/torch/jit/_trace.py", line 118, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1182, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/distilbert/modeling_distilbert.py", line 1225, in forward
    outputs = self.distilbert(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1182, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/distilbert/modeling_distilbert.py", line 806, in forward
    embeddings = self.embeddings(input_ids, inputs_embeds)  # (bs, seq_length, dim)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1182, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/distilbert/modeling_distilbert.py", line 129, in forward
    input_embeds = self.word_embeddings(input_ids)  # (bs, max_seq_length, dim)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1182, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 160, in forward
    return F.embedding(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.FloatTensor instead (while checking arguments for embedding)

[jfrery]

Hey my bad I missed your answer. Concrete ML doesn't support embedding layer. We will support that very soon.

If you don't mind, we can convert your issue to a Feature request to make sure we have the embedding supported asap.

[dopc]

Thanks for the answer and for converting the issue to a feature request.
Looking forward to the support.