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.