Fail to load NLP Bert Model in AWS Lambda Arm64 (Graviton 2)

Description

I'm testing nlp Bert performance using mxnet (including gluonnlp) in aws lambda based on arm64 (graviton2 processor). These scripts that I used

save Bert model with Gluonnlp

model_name_ = "bert_12_768_12"
dataset = "book_corpus_wiki_en_uncased"
model, _ = nlp.model.get_model(
                name=model_name_,
                dataset_name=dataset,
                pretrained=True,
                use_pooler=True,
                use_decoder=False,
                use_classifier=False,
            )
model = nlp.model.BERTClassifier(model, dropout=0.1, num_classes=2)
model.initialize(ctx=ctx)
model.hybridize(static_alloc=True)

mx_out = model(inputs_nd, token_types_nd, valid_length_nd)
mx_out.wait_to_read()

target_path = f"./{model_name}_{batch_size}"
from pathlib import Path
Path(target_path).mkdir(parents=True, exist_ok=True)  

model.export(f'{model_name}_{batch_size}/model')

load Bert in AWS Lambda based on ARM64

import mxnet as mx
import mxnet.ndarray as nd
from mxnet import nd, gluon

model_json, model_params = model_path + '/model-symbol.json', model_path + '/model-0000.params'
model = gluon.nn.SymbolBlock.imports(model_json, ['data0', 'data1', 'data2'], model_params, ctx=ctx)

Then, I faced this error

[ERROR] MXNetError: MXNetError: Invalid Parameter format for output_dim expect int but value='768', in operator Embedding(name="bertmodel0_word_embed_embedding0_fwd", sparse_grad="False", output_dim="768", dtype="float32", input_dim="30522")

I checked these scripts work without any issue in AWS Lambda based on X86 (Intel processor) and AWS EC2 Graviton2 environment.

Environment

AWS Lambda based on ARM64
python 3.8
mxnet 1.9.1

Could you help me to solve this issue? Thank you !

apache / mxnet

Fail to load NLP Bert Model in AWS Lambda Arm64 (Graviton 2) #21069

Description

save Bert model with Gluonnlp

load Bert in AWS Lambda based on ARM64

Environment