Add MegatronGPTDeployable for serving .nemo models in pytriton

[jukim-nv]

What does this PR do ?

We currently only support deploying TRTLLM models in the Triton inference server. This MR adds the necessary support to let us deploy .nemo models to Triton as well.

Collection: deploy

Changelog

• Adds MegatronLLMDeployable class which implements ITritonDeployable

Usage

• Usage is shown in the new tests/deploy/test_triton_deployable.py

    megatron_deployable = MegatronLLMDeployable(args.nemo_checkpoint, args.num_gpus)
    nm = DeployPyTriton(
        model=megatron_deployable,
        triton_model_name=model_name,
        triton_model_version=1,
        max_batch_size=8,
        port=8000,
        address="0.0.0.0",
        streaming=False,
    )
    nm.deploy()
    nm.run()

Jenkins CI

To run Jenkins, a NeMo User with write access must comment jenkins on the PR.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[oyilmaz-nvidia]

Can we add a test similar to this https://github.com/jukim-nv/NeMo/blob/megatrongpt_deployable/tests/export/test_nemo_export.py under tests/deploy/test_deploy_pytriton.py?

And in that test, can we test the class you created?

[oyilmaz-nvidia]

Let's change this name to PyTritonDeployableLLM and move this file under nemo/deploy/nlp

[jukim-nv]

Renamed to MegatronLLMDeployable for now. As discussed offline, whatever name we decide on for this class should match the Model class that we create inside of it. So if the final name is MegatronLLMDeployable, the internal self.model should be of type MegatronLLMModel.

[oyilmaz-nvidia]

Can we also add this pytorch level deployment in this script https://github.com/NVIDIA/NeMo/blob/main/scripts/deploy/nlp/deploy_triton.py please?

[oyilmaz-nvidia]

Can we please change the existing_model to model?

[jukim-nv]

Is there a particular convention we are trying to satisfy with renaming to model? I personally don't think that model alone is descriptive enough for the argument, considering that the word "model" can refer to many different things the context of AI and LLMs.
existing_model is not particularly good either so I would be open to other suggestions. Maybe preinitialized_model?

[oyilmaz-nvidia]

We basically pass a model in the memory. That's why I was suggesting the model. But I also agree that might be descriptive enough. maybe initialized_model?

[oyilmaz-nvidia]

Let's remove the according to EricH, but GPT-2B model will fail if it is enabled part and assume that it should be set to true for now.

[oyilmaz-nvidia]

let's use nemo logger.info here