Latent MOE support and fix KV cache quant export #768
Conversation
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #768 +/- ##
==========================================
- Coverage 74.62% 74.62% -0.01%
==========================================
Files 192 192
Lines 18989 18992 +3
==========================================
+ Hits 14171 14172 +1
- Misses 4818 4820 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
jenchen13
changed the title
| # Step 1: Sync amax across local experts in a SequentialMLP | ||
| for name, module in model.named_modules(): | ||
| if hasattr(module, "sync_moe_local_experts_amax"): | ||
| module.sync_moe_local_experts_amax() | ||
|
|
||
| # TODO just for testing | ||
| if "experts" in name and "weight_quantizer" in name: | ||
| assert child.amax is not None |
There was a problem hiding this comment.
Can we move this before if distributed sync check? This is not doing anything particular to distributed sync
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
| # TODO double check if MOE forward will be implemented in MoELayer or TransformerLayer | ||
| # We do not need both layers to be patched | ||
|
|
||
| @QuantModuleRegistry.register({megatron_transformer_layer.TransformerLayer: "megatron_transformer_layer_TransformerLayer"}) |
There was a problem hiding this comment.
TODO maybe remove this since MOE forward will not be removed in MLM main?
| "input_layernorm": NameRemapping("backbone.layers.{}.norm."), | ||
| "linear_qkv": QKVSlicing("backbone.layers.{}.mixer."), | ||
| "linear_proj": NameRemapping("backbone.layers.{}.mixer.o_proj."), | ||
| "core_attention": SelfAttentionScaling("backbone.layers.{}.mixer."), |
There was a problem hiding this comment.
doublecheck that this is only needed for export
Signed-off-by: jenchen13 <jennifchen@nvidia.com>
jenchen13
changed the title
2 participants
What does this PR do?
Type of change: New feature
Overview:
New nemotron models use
TransformerLayer.forward()instead ofMoELayer.forward()for MOE. This is a breaking change to our quantization implementation for Nano3 which relied on patchingMoELayer.forward()to force tokens to be routed to all experts during calibration.TransformerLayer.forward()which forces tokens to be routed to all experts during PTQ calibrationqkv_layer.output_quantizerexport & replace with properk/v_bmm_quantizerlogic? TODO Potentially remove MoELayer quant config if all MOEs in future will use TransformerLayer instead?
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Additional Information