## What does this PR do?
[Short term]: Megatron based tests take a long time often resulting in
CICD timeout. Splitting megatron tests into a dedicated CICD job for
faster overall CI/CD run
[Mid/Long term]: Run all megatron gpu tests using `torchrun` instead of
`pytest` so all dist processes are already created and all individual
tests no longer need to setup and destroy their processes which adds a
lot of overhead per test
## Testing
<!-- Mention how have you tested your change if applicable. -->
- [x] 1-GPU CI/CD passing (on this PR)
- [x] 2-GPU CI/CD passing (on nightly run - manually triggered):
https://github.com/NVIDIA/Model-Optimizer/actions/runs/22000517688
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
## What does this PR do?
**Type of change:** New feature and bug fix
**Overview:**
Support Latent MOE and Repeated MTP for NemotronH models
- Enable latent MOE modules during megatron import/export
- Fix KV cache quantization export: remove old
`qkv_layer.output_quantizer` export & replace with proper
`k/v_bmm_quantizer` logic
- Improvements to EP amax sync
- Support repeated MTP import/export for NemotronH models (only BF16
export for MTP for now)
## Usage
<!-- You can potentially add a usage example below. -->
```python
# Add a code snippet demonstrating how to use this
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
## 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/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **New Features**
* Added support for grouped MLP and self-attention scaling operations in
model export workflows
* Enhanced model parallel training capabilities with improved component
mapping
* Expanded quantization configuration handling with dynamic module
exclusion across distributed ranks
* Improved support for additional transformer engine components
* **Refactor**
* Reorganized internal export and import logic for improved
maintainability and specialist model architecture support
<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>
Signed-off-by: Jenny Chen <jennifchen@nvidia.com>
## What does this PR do?
**Type of change:** Fix MCore KV Cache Quantization: Amax Device
Placement Bug; Code clean up; Distributed Sync of KVCache Quantizer
params; unittest expansion to hybrid models
**Overview:** Fixes bugs preventing MCore KV Cache quantization from
working during checkpoint restore.
### Bug Chain
**Bug 1:** `is_enabled = self.weight_quantizer.is_enabled if
hasattr(self, "weight_quantizer") else False`
No `weight_quantizer` for KV-cache-only quant → `is_enabled=False` →
metadata not saved → `modelopt_post_restore()` never called. *(Thanks to
@jenchen13 )*
**Bug 2:** After fixing Bug 1, `_amax` restored on CPU (via
`_reset_pytorch_state_from_metadata`). Fallback
`_calibrate_quantizers()` never called because `_amax` exists.
**Bug 3:** Even if called, `_calibrate_quantizers()` fails —
`core_attention` has no parameters → can't determine device/dtype.
### The Fix
1. Remove `is_enabled` check entirely — disabled modules may still need
metadata restore. Explicitly skip `output_layer` from extra state
callbacks (never quantized)
2. Set `dtype`/`device` on `core_attention` from parent Attention
module, `modelopt_post_restore()` calls `self.to(device, dtype)`
3. Remove dead `_calibrate_quantizers()` code (will bring back similar
logic for KV cache affine quantization)
### Previous Unit Test Was Wrong
`model_test` was `mtq.quantize()`'d, not `mto.restore()`'d. Never tested
actual restore path.
### Additional Fixes
- Amax sync across DP/TP for KV cache quantizers
- `flash_decode` auto-disabled
### Code Cleanup
Removed ~100 lines of dead code.
## Testing
1. MCore KV Cache QAD with Nano V3 + Context Parallel works
2. Unit tests: hybrid models, KV+GEMM configs, correct restore workflow,
backward pass validation
## Before your PR is "*Ready for review*"
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: No
- **Did you update Changelog?**: Yes
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Asma Thekkumpate <akuriparambi@cw-dfw-cs-001-vscode-02.cm.cluster>
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>
- CI tests for min transformers version (4.48) started failing since new
peft 1.18 was released yesterday blocking all new PRs from merging
- Transformers v5.0 is right around the corner so we should not
recommend too older version
- Set new minimum transformers to 4.53 instead of 4.48
- nemo:25.07+ and trtllm:1.0.0+ docker containers have transformers
4.53+
- ModelOpt will still likely work for transformers 4.48-4.52 for now but
going forward will not be tested in CICD
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Keval Morabia