Repository navigation
Conversation
cosmos_framework.quantization imports ModelOpt, but no extra or dependency group declared it. Add a `quantization` extra pinning nvidia-modelopt==0.47.0 and regenerate uv.lock; it adds nvidia-modelopt and its dependency pulp, and leaves every existing locked package unchanged. Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
ModelOpt's generic fused-experts wrapper only intercepts F.linear, so on the framework's packed grouped-mm experts it quantizes nothing. - Add QuantPackedExperts, a ModelOpt wrapper for Qwen3VLMoeTextExpertsGroupedMm and MoEExperts with per-expert FP8 weight scales and one input scale per GEMM, and register_packed_experts() to install it before mtq.quantize. - Add assert_quantizers_calibrated() to fail on missing or invalid ranges. - Export packed experts in their source layout with [E,1,1] weight scales and scalar input scales, outside ModelOpt's linear-only exporter. - Skip the root Transformers modelopt_state.pth when the transformer has packed experts. Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
is_packed_expert_key is a plain key-name check, but export.py imported it from quantization/moe.py, which imports the MoE model code and Triton. That broke the CPU export tests, which load export.py without the package's __init__. Move it next to the other key-name rules in export.py, and import FP8_MAX in moe.py instead of duplicating it. Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
Yokto13
force-pushed
the
dfarhan/moe-fp8-experts
branch
from
October 27, 2026 11:00
27261b0 to
d4caba9
Compare
Document the steps the packed-expert path needs: register the wrapper before mtq.quantize or a restore, add per-expert weight quantizer entries to the config, and export with export_quantized_transformer. State that the experts are quantized as W8A8 with static input scales and that weight-only export is not supported. Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
Align the quantization dependency with the ModelOpt version used for Cosmos3 FP8 checkpoints.\n\nCommit cc28c0c removed unsupported nn.Embedding quantizer metadata to make checkpoints compatible with ModelOpt v0.44.0. Pin the dependency and lockfile to that version accordingly.\n\nContext: NVIDIA@cc28c0c
… calibrating weight quantizers QuantPackedExperts ran gate_up_proj_weight_quantizer and down_proj_weight_quantizer on the full 3D expert weight tensors on every forward pass during calibration, even though the weights are static and export_packed_experts quantizes them directly from their per-expert max. Remove the two weight quantizers from QuantPackedExperts and compute the per-expert weight amax directly from gate_up_proj and down_proj inside export_packed_experts. This avoids repeatedly reducing the expert weights during calibration and removes the need for callers to append axis=0 weight-quantizer rules to mtq.FP8_DEFAULT_CFG. Signed-off-by: Dominik Farhan <dfarhan@nvidia.com>
cosmos_framework.quantization.calibration imports `datasets` at module level, and the package __init__ imports calibration eagerly, so `import cosmos_framework.quantization` fails with ModuleNotFoundError when only the quantization extra is installed. `datasets` was only provided by the train extra, and nvidia-modelopt does not depend on it (it imports it lazily). Left unpinned to match the train extra; uv.lock keeps it at 4.6.1.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds code useful for quantizing packed MoE experts with ModelOpt and adds ModelOpt as extra dependency.
ModelOpt's fused-experts wrapper intercepts only
F.linearcalls and ignores thegrouped GEMMs (
torch._grouped_mm) used inQwen3VLMoeTextExpertsGroupedMmandMoEExperts, so it quantizes nothing there.Changes
that ModelOpt understands: it runs both grouped GEMMs itself, with ModelOpt
quantizers on their inputs and weights. Call
register_packed_experts()beforemtq.quantizeto use it.with per-expert weight scales.
quantizationextra topyproject.tomlwith ModelOpt(
nvidia-modelopt==0.47.0) and datasets package. ModelOpt is required bycosmos_framework.quantizationand its tests (currently quantization tests on main cannot be run because they import ModelOpt), but was not declared as a dependency before.
The changes touch only
quantization(and its tests) and the extras in dependencies.Dense models export as before.