added/fixed export for Megatron models

[Davood-M]

Signed-off-by: David Mosallanezhad [email protected]

What does this PR do ?

Fixed Megatron export for updated Megatron models

Collection: [Note which collection this PR will affect]

Changelog

• Changed tokens_head and row parallel layers to linear during the export process

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• [X] 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
• [X] 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)

[lgtm-com[bot]]

This pull request introduces 1 alert when merging d42d15cd85a14a6555276fa87fb22d89820d1fe8 into 80e8f9be8b8a4e6b831c4e403821ab2c7780f3b0 - view on LGTM.com

new alerts:

• 1 for Unused import

[lgtm-com[bot]]

This pull request introduces 1 alert when merging 785e48a8a53272eb8b15602e29a9f07bb93c18d0 into 80e8f9be8b8a4e6b831c4e403821ab2c7780f3b0 - view on LGTM.com

new alerts:

• 1 for Unused import

[lgtm-com[bot]]

This pull request introduces 1 alert when merging 1a0f1e3c47d8a1edb3cc029dd55eb58e1d82665b into a309c3d80a10f07f5939e37a82f31b3ee30f7124 - view on LGTM.com

new alerts:

• 1 for Unused import