mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-10 09:34:01 +08:00
kinjalpatel27
f7313790f9
iCo-authored-by: Asma Kuriparambil Thekkumpate <akuriparambi@nvidia.com> ## What does this PR do? **Type of change:** New Feature **Overview:** This MR adds support for quantizing TE Ops in megatron, specifically TERowParallelLinear, TEColumnParallelLinear and TELayerNormColumnParallelLinear. ## Usage It can be used by enabling TE spec in megatron ## Testing Added unit tests for testing functionality `test_homogeneous_sharded_state_dict_te_spec` `test_convert_mcore_te_gpt_model` `test_quantize_forward_backward` ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes ## Additional Information <!-- E.g. related issue. --> --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>