trtllm-build ignores `--model_cls_file` and `--model_cls_name`

[abhishekudupa]

The trtllm-build entrypoint ignores the user-provided --model_cls_file and --model_cls_name. It looks like although the model class is resolved and imported, an appropriate entry in the MODEL_MAP isn't made, as seen here.

I've verified that adding the line MODEL_MAP[args.model_cls_name] = model_cls immediately after the line referenced above works well to fix this issue.

Who can help?

@ncomly-nvidia , @byshiue

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Try a trtllm-build command with any model architecture that is not supported by default, by using the --model_cls_file and --model_cls_name, you'll end up with a KeyError.

Expected behavior

The trtllm-build picks up my model classes and runs to completion and produces the engine files.

actual behavior

Raises a KeyError

additional notes

None, fix proposed in problem description.

[abhishekudupa]

Team, any updates on when I can expect the fix? Is the fix included in trtllm-0.15? Thx!

[nv-guomingz]

Hi @QiJune would u please take a look this question?