diff --git a/cosmos_framework/quantization/export.py b/cosmos_framework/quantization/export.py index 62485743e..b9de291cb 100644 --- a/cosmos_framework/quantization/export.py +++ b/cosmos_framework/quantization/export.py @@ -98,6 +98,11 @@ def remap_framework_state_dict( for suffix in _SCALE_SUFFIXES: if not net_key.endswith(suffix): continue + # Packed MoE expert scales sit beside the packed parameter itself. + param_diffusers_key = net_to_diffusers.get(net_key[: -len(suffix)]) + if param_diffusers_key is not None: + diffusers_key = param_diffusers_key + suffix + break weight_key = net_key[: -len(suffix)] + ".weight" weight_diffusers_key = net_to_diffusers.get(weight_key) if weight_diffusers_key is not None: @@ -111,6 +116,22 @@ def remap_framework_state_dict( return remapped +def _remap_packed_experts( + experts: dict[str, torch.Tensor], bf16_state_dict: dict[str, torch.Tensor] +) -> dict[str, torch.Tensor]: + """:func:`remap_framework_state_dict` for packed experts, where every tensor must map. + + A dropped expert would leave its BF16 source tensor without scales under an FP8 config. + """ + remapped = remap_framework_state_dict(experts, bf16_state_dict) + if len(remapped) != len(experts): + unmapped = sorted( + key for key in experts if not remap_framework_state_dict({key: experts[key]}, bf16_state_dict) + ) + raise ValueError(f"{len(unmapped)} packed expert tensor(s) have no Diffusers key: {unmapped[:10]}") + return remapped + + # Top-level projections produced by reverse_remap use the cosmos3_vfm names; the # bf16 diffusers base provides the equivalently-named proj_in/proj_out, so we drop # ours and keep the base copies. @@ -118,6 +139,17 @@ def remap_framework_state_dict( _FP8_DTYPES = {torch.float8_e4m3fn, getattr(torch, "float8_e5m2", torch.float8_e4m3fn)} FP8_MAX = 448.0 _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale", ".input_scale_2", ".alpha", ".pre_quant_scale") +_PACKED_EXPERT_KEY_SUFFIXES = tuple( + f"experts.{name}{suffix}" + for name in ("gate_up_proj", "down_proj") + for suffix in ("", ".weight_scale", ".input_scale") +) + + +def is_packed_expert_key(key: str) -> bool: + """True for a packed expert parameter or its scale, e.g. ``...experts.gate_up_proj.weight_scale``.""" + return key.endswith(_PACKED_EXPERT_KEY_SUFFIXES) + # Linears the diffusers checkpoint keeps in bf16 (never quantized), emitted in the # ModelOpt ``ignore`` list so a reasoning-path loader (vllm Cosmos3ForConditionalGeneration) @@ -207,7 +239,12 @@ def _remap_framework_modelopt_state( def remap_module(name: str) -> str | None: return module_map.get(name.removeprefix("net.")) - active_modules = {key.removesuffix(".weight_scale") for key in exported_state_dict if key.endswith(".weight_scale")} + # Packed MoE experts are not linears and have no component graph state. + active_modules = { + key.removesuffix(".weight_scale") + for key in exported_state_dict + if key.endswith(".weight_scale") and not is_packed_expert_key(key) + } real_quantize = modes["real_quantize"] real_metadata = real_quantize["metadata"] @@ -381,7 +418,7 @@ def overlay(self, bf16: dict, quantized: dict) -> dict: merged[k] = v n_scale += 1 continue - if k.endswith(".weight") and v.dtype in _FP8_DTYPES: + if (k.endswith(".weight") or is_packed_expert_key(k)) and v.dtype in _FP8_DTYPES: if k not in bf16: raise KeyError( f"quantized weight {k!r} has no bf16 counterpart — layer-key " @@ -390,7 +427,9 @@ def overlay(self, bf16: dict, quantized: dict) -> dict: if bf16[k].shape != v.shape: raise ValueError(f"shape mismatch for {k}: bf16 {tuple(bf16[k].shape)} vs fp8 {tuple(v.shape)}") merged[k] = v - quantized_modules.add(k[: -len(".weight")]) + # Packed experts keep only their scales; the quantizer buffers below are per linear. + if not is_packed_expert_key(k): + quantized_modules.add(k[: -len(".weight")]) n_weight += 1 continue # Remaining bf16 tensors (norms, embed_tokens, renamed projections) are @@ -423,6 +462,7 @@ def export( modelopt_config: dict | None = None, ) -> None: from .checkpoint_io import load_sharded_safetensors, save_sharded_safetensors + from .moe import QuantPackedExperts, export_packed_experts net = mdl.net transformer_export_dir.mkdir(parents=True, exist_ok=True) @@ -431,6 +471,14 @@ def export( if n_collapsed > 0: print(f"[export] collapsed per-channel input amax to scalar on {n_collapsed} quantizers") + # Packed MoE experts are exported directly; ModelOpt's compress and unified + # exporter only handle linears, so their quantizers are disabled afterwards. + experts = export_packed_experts(net) + for module in net.modules(): + if isinstance(module, QuantPackedExperts): + for name in ("gate_up_proj", "down_proj"): + getattr(module, f"{name}_input_quantizer").disable() + # Record ModelOpt's real-quantized graph metadata before the unified # exporter unpacks those weights into plain FP8 checkpoint tensors. mtq.compress(mdl) @@ -446,6 +494,7 @@ def export( with hide_quantizers_from_state_dict(net): net_state_dict = {k: v.detach().contiguous().cpu() for k, v in net.state_dict().items()} quantized = remap_framework_state_dict(net_state_dict, bf16) + quantized.update(_remap_packed_experts(experts, bf16)) merged = self.overlay(bf16, quantized) n_shards = save_sharded_safetensors(merged, transformer_export_dir) print(f"[export] wrote {len(merged)} tensors across {n_shards} shard(s)") @@ -505,12 +554,13 @@ def _write_root_hf_quant_config(root_dir: Path, transformer_dir: Path) -> None: print(f"[assemble] wrote {root_dir / 'hf_quant_config.json'} for upstream-vLLM discovery") -def _write_root_weight_index(root_dir: Path) -> None: - """(Re)generate the root ``model.safetensors.index.json`` for whole-model loading.""" +def _write_root_weight_index(root_dir: Path) -> dict[str, str]: + """(Re)generate the root ``model.safetensors.index.json``; return its weight map.""" index = build_root_index(root_dir) dst = root_dir / "model.safetensors.index.json" index.write(dst) print(f"[assemble] wrote {dst.name} ({len(index.weight_map)} tensors, {index.metadata['total_size']} bytes)") + return index.weight_map def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer_dir: Path): @@ -520,6 +570,7 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer vision_encoder symlinked back); only ``transformer/`` is physically new (the quantized export). A fresh root weight index + a root ``hf_quant_config.json`` are written so both the diffusers and transformers loaders resolve the new weights. + The root Transformers ModelOpt state is skipped when the transformer has packed experts. """ input_dir, output_dir = Path(input_dir), Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) @@ -552,8 +603,12 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer rel_src = os.path.relpath(src, start=output_dir.resolve()) print(f"[assemble] linking {entry.name} -> {rel_src}") dst.symlink_to(rel_src) - _write_root_weight_index(output_dir) - component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False) - Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir) + weight_map = _write_root_weight_index(output_dir) + has_packed_experts = any( + is_packed_expert_key(key) for key, file in weight_map.items() if file.startswith("transformer/") + ) + if not has_packed_experts: + component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False) + Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir) _write_root_hf_quant_config(output_dir, target_transformer) print(f"[assemble] drop-in dir ready at {output_dir}") diff --git a/cosmos_framework/quantization/moe.py b/cosmos_framework/quantization/moe.py new file mode 100644 index 000000000..a711c26f7 --- /dev/null +++ b/cosmos_framework/quantization/moe.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""ModelOpt FP8 quantization of the packed grouped-mm MoE experts. + +Calibrates and fake-quantizes the expert inputs through the wrapper's ``forward``, and quantizes +the expert weights (with one scale per expert) directly in :func:`export_packed_experts`. +This code does not support weight-only. + +How to run quantization: call :func:`register_packed_experts` before ``mtq.quantize`` (and before +restoring a saved ModelOpt state), then export with +:func:`~cosmos_framework.quantization.export.export_quantized_transformer`, which writes the +experts via :func:`export_packed_experts`. + +Without the registration, ModelOpt's generic ``_QuantFusedExperts`` claims these experts and, +since it only intercepts ``F.linear`` and they run ``torch._grouped_mm``, quantizes nothing. +""" + +from __future__ import annotations + +import torch +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer +from torch import nn + +from cosmos_framework.model.generator.moe.interface import MoEExperts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import Qwen3VLMoeTextExpertsGroupedMm +from cosmos_framework.quantization.export import FP8_MAX + +_PACKED_EXPERT_TYPES = (Qwen3VLMoeTextExpertsGroupedMm, MoEExperts) + + +class QuantPackedExperts(QuantModule): + """Calibrates and fake-quantizes expert inputs; per-expert FP8 weights are quantized on export.""" + + def _setup(self): + self.gate_up_proj_input_quantizer = TensorQuantizer() + self.down_proj_input_quantizer = TensorQuantizer() + + def forward( + self, + hidden_states: torch.Tensor, # [num_tokens,hidden_size] + topk_scores: torch.Tensor, # [num_tokens,top_k] + expert_indices: torch.Tensor, # [num_tokens,top_k] + num_tokens_per_expert: torch.Tensor, # [num_experts] + token_mask: torch.Tensor | None = None, # [num_tokens] + ) -> torch.Tensor: # [num_tokens,hidden_size] + # Group the real (token, slot) pairs by expert; ``num_tokens_per_expert`` counts exactly these. + top_k = expert_indices.shape[1] + slots = torch.arange(expert_indices.numel(), device=expert_indices.device) # [num_tokens*top_k] + if token_mask is not None: + slots = slots[token_mask.repeat_interleave(top_k)] # [num_routed] + if slots.numel() == 0: + return torch.zeros_like(hidden_states) # [num_tokens,hidden_size] + slots = slots[torch.argsort(expert_indices.reshape(-1)[slots], stable=True)] # [num_routed] + tokens = slots // top_k # [num_routed] + scores = topk_scores.reshape(-1)[slots].to(hidden_states.dtype) # [num_routed] + x = hidden_states[tokens] # [num_routed,hidden_size] + routed = self._quantized_experts(x, num_tokens_per_expert, scores) # [num_routed,hidden_size] + return torch.zeros_like(hidden_states).index_add_(0, tokens, routed) # [num_tokens,hidden_size] + + def _quantized_experts( + self, + x: torch.Tensor, # [num_routed,hidden_size] (grouped by expert) + num_tokens_per_expert: torch.Tensor, # [num_experts] + scores: torch.Tensor, # [num_routed] + ) -> torch.Tensor: # [num_routed,hidden_size] + # a2 observes ``act(gate) * up`` before the routing scores, because vLLM applies the scores + # after the down GEMM. Per-expert weight scales are computed directly on export. + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) # [num_experts] + x = self.gate_up_proj_input_quantizer(x) + h = torch._grouped_mm(x, self.gate_up_proj, offs=offsets) + g, u = torch.chunk(h, chunks=2, dim=-1) # 2x [num_routed,moe_intermediate_size] + a = self.down_proj_input_quantizer(self.act_fn(g) * u) # [num_routed,moe_intermediate_size] + return torch._grouped_mm(a * scores.unsqueeze(-1), self.down_proj, offs=offsets) # [num_routed,hidden_size] + + +def register_packed_experts() -> None: + """Register :class:`QuantPackedExperts` before ``mtq.quantize`` so ``_QuantFusedExperts`` skips it.""" + pending: dict[type[nn.Module], str] = {} + for expert_cls in _PACKED_EXPERT_TYPES: + registered = QuantModuleRegistry.get(expert_cls) + if registered is not None and issubclass(registered, QuantPackedExperts): + continue + if registered is not None: + # An earlier ``mtq.quantize`` that saw this type unregistered has globally registered + # ModelOpt's generic ``_QuantFusedExperts`` for it, which would quantize nothing here. + QuantModuleRegistry.unregister(expert_cls) + pending[expert_cls] = expert_cls.__name__ + if pending: + QuantModuleRegistry.register(pending)(QuantPackedExperts) + + +def has_packed_experts(model: nn.Module) -> bool: + return any(isinstance(module, _PACKED_EXPERT_TYPES) for module in model.modules()) + + +def assert_quantizers_calibrated(model: nn.Module) -> None: + """Fail export instead of writing a missing or zero scale for an enabled quantizer.""" + problems = [] + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and not module.is_enabled: + continue + if isinstance(module, TensorQuantizer): + amax = getattr(module, "amax", None) + if amax is None: + problems.append(f"{name}: no amax") + elif not torch.isfinite(amax).all(): + problems.append(f"{name}: non-finite amax") + elif name.endswith("input_quantizer") and not (amax > 0).all(): + problems.append(f"{name}: non-positive amax") + elif ( + isinstance(getattr(module, "gate_up_proj", None), nn.Parameter) + and module.gate_up_proj.ndim == 3 + and not isinstance(module, QuantPackedExperts) + ): + problems.append(f"{name}: packed experts {type(module).__name__} are not quantized") + if problems: + shown = "; ".join(problems[:10]) + raise ValueError(f"{len(problems)} quantizer problem(s), first {min(len(problems), 10)}: {shown}") + + +_PACKED_EXPERT_QUANTIZERS = ( + "gate_up_proj_input_quantizer", + "down_proj_input_quantizer", +) + + +def _check_per_expert_fp8(name: str, module: QuantPackedExperts) -> None: + """Refuse any input quantizer layout but FP8 E4M3 with a scalar amax.""" + for quantizer_name in _PACKED_EXPERT_QUANTIZERS: + quantizer = getattr(module, quantizer_name) + amax_shape = tuple(quantizer.amax.shape) + if quantizer.num_bits != (4, 3) or quantizer.block_sizes: + found = f"num_bits={quantizer.num_bits}, block_sizes={quantizer.block_sizes}" + elif quantizer.amax.numel() == 1: + continue + else: + found = f"amax shape {amax_shape}, expected a single element" + raise ValueError(f"{name}.{quantizer_name}: {found}; packed-expert export supports per-expert FP8 only") + + +def export_packed_experts(net: nn.Module) -> dict[str, torch.Tensor]: + """FP8 packed experts in the source layout, with per-expert weight scales and scalar input scales.""" + if not has_packed_experts(net): + return {} # Dense models export exactly as before, without the calibration check. + assert_quantizers_calibrated(net) + exported = {} + for name, module in net.named_modules(): + if not isinstance(module, QuantPackedExperts): + continue + enabled = [getattr(module, quantizer).is_enabled for quantizer in _PACKED_EXPERT_QUANTIZERS] + if not any(enabled): + continue # Skipped by the coverage filter; the experts stay BF16 from the source. + if not all(enabled): + raise ValueError(f"{name or 'experts'}: packed-expert quantizers are only partly enabled: {enabled}") + _check_per_expert_fp8(name or "experts", module) + for param_name in ("gate_up_proj", "down_proj"): + # BF16 models keep BF16 amax; the scales are FP32 on disk. + param = getattr(module, param_name).detach() + weight_amax = param.abs().amax(dim=(1, 2), keepdim=True).float() + weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / FP8_MAX) # [num_experts,1,1] + fp8 = (param.float() / weight_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + input_amax = getattr(module, f"{param_name}_input_quantizer").amax.detach().float() + key = f"{name}.{param_name}" if name else param_name + exported[key] = fp8.contiguous().cpu() + exported[f"{key}.weight_scale"] = weight_scale.contiguous().cpu() + exported[f"{key}.input_scale"] = (input_amax / FP8_MAX).reshape(()).contiguous().cpu() + return exported diff --git a/pyproject.toml b/pyproject.toml index 386421dba..2568f0b32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ [project.optional-dependencies] automl = ["nvidia-automl-core==0.1.0"] automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"] +quantization = ["datasets", "nvidia-modelopt==0.44.0"] reasoner-eval = [ "attrs", "evaluate", diff --git a/tests/quantization/modelopt_state_test.py b/tests/quantization/modelopt_state_test.py index 718b8830a..a04c31f6c 100644 --- a/tests/quantization/modelopt_state_test.py +++ b/tests/quantization/modelopt_state_test.py @@ -365,6 +365,68 @@ def test_assembly_keeps_distinct_states_and_source_config(export_modules, tmp_pa } +def test_assemble_skips_root_modelopt_state_for_packed_experts(export_modules, tmp_path, monkeypatch): + _, exporter = export_modules + source, output, staging = (tmp_path / name for name in ("source", "output", "staging")) + for directory in (source, output, staging): + directory.mkdir() + (staging / "config.json").write_text(json.dumps({"quantization_config": {"quant_method": "modelopt"}})) + save_file({"layers.0.mlp.experts.gate_up_proj": torch.zeros(2, 4, 4)}, str(staging / "model.safetensors")) + torch.save({}, staging / "modelopt_state.pth") + + def fail(*args, **kwargs): + raise AssertionError("root ModelOpt state must not be built") + + monkeypatch.setattr(exporter, "build_reasoner_modelopt_state", fail) + exporter.assemble_output_dir(source, output, staging) + assert not (output / "modelopt_state.pth").exists() + assert (output / "hf_quant_config.json").is_file() + + +def test_component_state_ignores_packed_experts(export_modules): + from cosmos_framework.inference.model import _diffusers_to_net_key + + _, exporter = export_modules + linear, experts = "layers.0.self_attn.to_q", "layers.0.mlp.experts.gate_up_proj" + net_linear = _diffusers_to_net_key(f"{linear}.weight").removesuffix(".weight") + model = torch.nn.Module() + model.add_module("net", torch.nn.Module()) + parent = model.net + *parents, leaf = net_linear.split(".") + for part in parents: + parent.add_module(part, torch.nn.Module()) + parent = parent._modules[part] + parent.add_module(leaf, torch.nn.Linear(16, 16, bias=False).bfloat16()) + model = apply_mode(model, [("quantize", copy.deepcopy(mtq.FP8_DEFAULT_CFG))]) + module = model.get_submodule(f"net.{net_linear}") + module.input_quantizer.amax = torch.tensor(3.5) + module.weight_quantizer.amax = module.weight.detach().float().abs().max() + mtq.compress(model) + bf16 = {f"{linear}.weight": torch.zeros(16, 16, dtype=torch.bfloat16), experts: torch.zeros(4, 16, 16)} + exported = { + f"{linear}.weight": torch.zeros(16, 16, dtype=torch.float8_e4m3fn), + experts: torch.zeros(4, 16, 16, dtype=torch.float8_e4m3fn), + } + for key in (linear, experts): + exported[f"{key}.weight_scale"] = torch.ones(()) + exported[f"{key}.input_scale"] = torch.ones(()) + + state = exporter._remap_framework_modelopt_state( + mto.modelopt_state(model), + bf16_state_dict=bf16, + exported_state_dict=exported, + modelopt_config=copy.deepcopy(mtq.FP8_DEFAULT_CFG), + ) + + modes = dict(state["modelopt_state_dict"]) + assert set(modes["real_quantize"]["metadata"]["q_tensor_state"]) == {linear} + names = [ + *modes["quantize"]["metadata"]["quantizer_state"], + *modes["real_quantize"]["metadata"]["real_quantizer_state"], + ] + assert names and not any("experts" in name for name in names) + + @pytest.mark.skipif(not os.environ.get("COSMOS3_REASONER_STATE_SOT"), reason="external SOT path not supplied") def test_matches_external_sot_metadata(export_modules, tmp_path): helper, _ = export_modules diff --git a/tests/quantization/moe_experts_test.py b/tests/quantization/moe_experts_test.py new file mode 100644 index 000000000..73e15f303 --- /dev/null +++ b/tests/quantization/moe_experts_test.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""ModelOpt FP8 calibration and export of the packed grouped-mm MoE experts. + +The CPU tests replace ``torch._grouped_mm`` and Triton's ``_generate_permute_indices`` with +PyTorch references. +""" + +import copy + +import modelopt.torch.quantization as mtq +import pytest +import torch +import torch.nn.functional as F +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import _QuantFusedExperts +from torch import nn + +import cosmos_framework.model.generator.moe.native as native_grouped_mm +from cosmos_framework.inference.model import _diffusers_to_net_key +from cosmos_framework.model.generator.moe.interface import MoEExperts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeTextConfig +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import ( + Qwen3VLMoeTextExpertsGroupedMm, + Qwen3VLMoeTextExpertsNaive, +) +from cosmos_framework.quantization.export import Fp8DiffusersExporter, _remap_packed_experts, remap_framework_state_dict +from cosmos_framework.quantization.moe import ( + QuantPackedExperts, + assert_quantizers_calibrated, + export_packed_experts, + register_packed_experts, +) + +QUANTIZERS = ( + "gate_up_proj_input_quantizer", + "down_proj_input_quantizer", +) + + +def _reference_grouped_mm(x: torch.Tensor, w: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + # torch._grouped_mm leaves the rows past offs[-1] uninitialized; a large value stands in for that garbage. + out = x.new_full((x.shape[0], w.shape[-1]), 1e4) + start = 0 + for group, end in enumerate(offs.tolist()): + out[start:end] = x[start:end] @ w[group] + start = end + return out + + +def _reference_permute_indices( + num_tokens_per_expert: torch.Tensor, num_experts: int, max_len: int, alignment: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch version of the Triton ``_generate_permute_indices``: -1 marks a padding slot.""" + sizes = num_tokens_per_expert.clamp_min(alignment).to(torch.int32) + sizes = (sizes + alignment - 1) // alignment * alignment # [num_experts] + permuted_indices = torch.full((max_len,), -1, dtype=torch.int32) + start = write = 0 + for count, size in zip(num_tokens_per_expert.tolist(), sizes.tolist(), strict=True): + permuted_indices[write : write + count] = torch.arange(start, start + count, dtype=torch.int32) + start, write = start + count, write + size + return permuted_indices, sizes + + +@pytest.fixture +def reference_routing(monkeypatch): + monkeypatch.setattr(torch, "_grouped_mm", _reference_grouped_mm) + monkeypatch.setattr(native_grouped_mm, "_generate_permute_indices", _reference_permute_indices) + + +def _config() -> Qwen3VLMoeTextConfig: + return Qwen3VLMoeTextConfig( + num_experts=4, hidden_size=16, moe_intermediate_size=8, num_experts_per_tok=2, hidden_act="silu" + ) + + +class Model(nn.Module): + def __init__(self, experts: nn.Module | None = None): + super().__init__() + self.experts = Qwen3VLMoeTextExpertsGroupedMm(_config()) if experts is None else experts + generator = torch.Generator().manual_seed(0) + for param in self.experts.parameters(): + param.data.normal_(generator=generator) + + +# The packed expert types ``register_packed_experts`` claims, built for a given top_k. +EXPERTS = { + "qwen_grouped_mm": lambda top_k: Qwen3VLMoeTextExpertsGroupedMm(_config(), top_k=top_k), + "moe_experts": lambda top_k: MoEExperts(hidden_size=16, intermediate_size=8, num_experts=4, top_k=top_k), +} + + +def _forward_inputs( + seed: int = 3, dtype: torch.dtype = torch.float32, top_k: int = 2, routing: str = "padding" +) -> tuple: + """``forward`` arguments for 6 tokens, top-``top_k`` of 4 experts, padded per ``routing``. + + Only the last token picks expert 2, so padding it leaves expert 2 without tokens. + """ + generator = torch.Generator().manual_seed(seed) + hidden_states = torch.randn(6, 16, generator=generator, dtype=dtype) + topk_scores = torch.rand(6, top_k, generator=generator, dtype=dtype) + expert_indices = torch.tensor([[0, 1, 3], [1, 3, 0], [0, 3, 1], [3, 0, 1], [1, 0, 3], [2, 1, 0]]) + expert_indices = expert_indices[:, :top_k].contiguous() + token_mask = torch.tensor([True, True, True, True, True, False]) + if routing == "no_padding": + token_mask = None + elif routing == "all_padding": + token_mask = torch.zeros(6, dtype=torch.bool) + elif routing == "empty": + hidden_states, topk_scores = hidden_states[:0], topk_scores[:0] + expert_indices, token_mask = expert_indices[:0], token_mask[:0] + else: + assert routing == "padding", routing + real = expert_indices if token_mask is None else expert_indices[token_mask] + num_tokens_per_expert = torch.bincount(real.reshape(-1), minlength=4) + return hidden_states, topk_scores, expert_indices, num_tokens_per_expert, token_mask + + +def _quantize(model: nn.Module, calibration: list[tuple] | None = None, *extra_cfg: dict) -> nn.Module: + register_packed_experts() + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"] = [*config["quant_cfg"], *extra_cfg] + + def forward_loop(model: nn.Module) -> None: + for inputs in calibration or [_forward_inputs()]: + model.experts(*inputs) + if hasattr(model, "linear"): + model.linear(inputs[0]) + + return mtq.quantize(model, config, forward_loop) + + +@pytest.fixture +def unregistered_experts(): + if QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm) is not None: + QuantModuleRegistry.unregister(Qwen3VLMoeTextExpertsGroupedMm) + yield + register_packed_experts() + + +@pytest.mark.usefixtures("reference_routing", "unregistered_experts") +def test_registration_replaces_generic_fused_experts_wrapper(): + generic = Model() + mtq.quantize(generic, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda model: None) + assert isinstance(generic.experts, _QuantFusedExperts) + + model = _quantize(Model()) + + assert isinstance(model.experts, QuantPackedExperts) + assert issubclass(QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm), QuantPackedExperts) + + +@pytest.mark.usefixtures("reference_routing") +@pytest.mark.parametrize("routing", ["padding", "no_padding", "all_padding", "empty"]) +@pytest.mark.parametrize("top_k", [2, 3]) +@pytest.mark.parametrize("experts", sorted(EXPERTS)) +def test_wrapper_matches_native_with_quantizers_disabled(experts, top_k, routing): + # QuantPackedExperts has its own forward; this catches it drifting from the native one. + model = Model(EXPERTS[experts](top_k)) + native = copy.deepcopy(model.experts) + inputs = _forward_inputs(seed=4, top_k=top_k, routing=routing) + expected = native(*inputs) + model = _quantize(model, [_forward_inputs(top_k=top_k)]) + assert isinstance(model.experts, QuantPackedExperts) + + for name in QUANTIZERS: + getattr(model.experts, name).disable() + # On CPU ``MoEExperts`` runs the dense ``native.naive``, so the two match up to rounding only. + torch.testing.assert_close(model.experts(*inputs), expected, rtol=1e-5, atol=1e-3) + + +@pytest.mark.usefixtures("reference_routing") +def test_each_input_quantizer_changes_the_wrapper_output(): + model = _quantize(Model()) + inputs = _forward_inputs(seed=4) + for name in QUANTIZERS: + getattr(model.experts, name).disable() + unquantized = model.experts(*inputs) + + for name in QUANTIZERS: + quantizer = getattr(model.experts, name) + quantizer.enable() + assert not torch.equal(model.experts(*inputs), unquantized), name + quantizer.disable() + + +@pytest.mark.usefixtures("reference_routing") +def test_forward_calibrates_unpadded_activations(): + calibration = [] + for seed in (1, 2): + hidden_states, _, expert_indices, counts, token_mask = _forward_inputs(seed) + hidden_states[~token_mask] = 1e3 # A masked token must not reach a1. + calibration.append((hidden_states, torch.full((6, 2), 0.25), expert_indices, counts, token_mask)) + + model = _quantize(Model(), calibration) + + assert_quantizers_calibrated(model) + experts = model.experts + real = [(x[mask], indices[mask]) for x, _, indices, _, mask in calibration] + assert torch.equal(experts.gate_up_proj_input_quantizer.amax, torch.stack([x.abs().max() for x, _ in real]).max()) + # a2 sees act(gate) * up on the real (token, expert) rows only, before the routing scores. + gate_up = torch.cat([torch.einsum("th,tkhf->tkf", x, experts.gate_up_proj[indices]) for x, indices in real]) + gate, up = gate_up.chunk(2, dim=-1) + expected = (F.silu(gate) * up).abs().max() + amax = experts.down_proj_input_quantizer.amax + torch.testing.assert_close(amax, expected) + assert not torch.allclose(amax, 0.25 * expected) + + +def test_assert_quantizers_calibrated_rejects_uncalibrated_and_unwrapped(): + model = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + with pytest.raises(ValueError, match=r"1\.input_quantizer"): + assert_quantizers_calibrated(model) + + model = nn.Sequential(nn.Linear(4, 4)) + model.experts = Qwen3VLMoeTextExpertsNaive(_config()) + # ModelOpt calibrates the generic fused-experts weights, so they must not be uninitialized memory. + generator = torch.Generator().manual_seed(0) + for param in model.experts.parameters(): + param.data.normal_(generator=generator) + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + with pytest.raises(ValueError, match="experts"): + assert_quantizers_calibrated(model) + + +@pytest.mark.usefixtures("reference_routing") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_export_packed_experts_layout(dtype): + model = Model().to(dtype) + model.experts.gate_up_proj.data[1] = 0 + model.experts.down_proj.data[1] = 0 + model = _quantize(model, [_forward_inputs(dtype=dtype)]) + + exported = export_packed_experts(model) + + shapes = {"gate_up_proj": (4, 16, 16), "down_proj": (4, 8, 16)} + assert set(exported) == { + f"experts.{name}{suffix}" for name in shapes for suffix in ("", ".weight_scale", ".input_scale") + } + for name, shape in shapes.items(): + weight = getattr(model.experts, name).float() + fp8 = exported[f"experts.{name}"] + weight_scale = exported[f"experts.{name}.weight_scale"] + input_scale = exported[f"experts.{name}.input_scale"] + assert fp8.dtype == torch.float8_e4m3fn and fp8.shape == shape + assert weight_scale.dtype == torch.float32 and weight_scale.shape == (4, 1, 1) + assert input_scale.dtype == torch.float32 and input_scale.shape == () + weight_amax = weight.abs().amax(dim=(1, 2), keepdim=True) + torch.testing.assert_close(weight_scale, torch.where(weight_amax == 0, 1.0, weight_amax / 448)) + assert weight_scale[1].item() == 1.0 # The all-zero expert gets a unit scale, not 0. + assert torch.isfinite(fp8[1].float()).all() and (fp8[1].float() == 0).all() + amax = getattr(model.experts, f"{name}_input_quantizer").amax.float() + torch.testing.assert_close(input_scale, amax / 448) + # 32 is the FP8 E4M3 step in the top binade [256, 448]. + assert ((fp8.float() * weight_scale - weight).abs() <= 32 * weight_scale).all() + assert all( + tensor.is_contiguous() and tensor.device.type == "cpu" for tensor in (fp8, weight_scale, input_scale) + ) + + +@pytest.mark.usefixtures("reference_routing") +def test_export_skips_dense_and_disabled_experts_and_rejects_mixed(): + dense = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + mtq.quantize(dense, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + assert export_packed_experts(dense) == {} # Uncalibrated, but dense models skip the check. + + # A coverage skip disables the experts before calibration, so they never get an amax. The + # linear keeps a weight to quantize: ModelOpt (>= 0.47) refuses a config with none enabled. + skipped = Model() + skipped.linear = nn.Linear(16, 16) + skipped = _quantize(skipped, None, {"quantizer_name": "*experts*", "enable": False}) + assert isinstance(skipped.experts, QuantPackedExperts) + assert export_packed_experts(skipped) == {} + + mixed = _quantize(Model()) + mixed.experts.down_proj_input_quantizer.disable() + with pytest.raises(ValueError, match="experts"): + export_packed_experts(mixed) + + +@pytest.mark.usefixtures("reference_routing") +def test_export_rejects_non_per_expert_fp8(): + per_channel_inputs = _quantize(Model(), None, {"quantizer_name": "*input_quantizer", "cfg": {"axis": 0}}) + assert per_channel_inputs.experts.gate_up_proj_input_quantizer.amax.numel() > 1 + with pytest.raises(ValueError, match="per-expert FP8") as error: + export_packed_experts(per_channel_inputs) + assert "experts" in str(error.value) and "gate_up_proj_input_quantizer" in str(error.value) + + int8 = _quantize(Model(), None, {"quantizer_name": "*experts*", "cfg": {"num_bits": 8}}) + assert int8.experts.down_proj_input_quantizer.num_bits == 8 + with pytest.raises(ValueError, match="per-expert FP8") as error: + export_packed_experts(int8) + assert "num_bits=8" in str(error.value) + + +def test_remap_packed_expert_scales(): + prefix = "layers.0.mlp.experts" + bf16 = {f"{prefix}.gate_up_proj": torch.zeros(4, 16, 16, dtype=torch.bfloat16)} + net_prefix = _diffusers_to_net_key(f"{prefix}.gate_up_proj", "transformer/weights.safetensors") + exported = {net_prefix + suffix: torch.ones(1) for suffix in ("", ".weight_scale", ".input_scale")} + expected = {f"{prefix}.gate_up_proj{suffix}" for suffix in ("", ".weight_scale", ".input_scale")} + + assert set(remap_framework_state_dict(exported, bf16)) == expected + assert set(_remap_packed_experts(exported, bf16)) == expected + exported["net.unknown.experts.down_proj.weight_scale"] = torch.ones(1) + with pytest.raises(ValueError, match=r"net\.unknown\.experts\.down_proj\.weight_scale"): + _remap_packed_experts(exported, bf16) + + +@pytest.mark.usefixtures("reference_routing") +def test_overlay_accepts_packed_fp8_experts(): + model = _quantize(Model()) + bf16 = { + f"layers.0.mlp.{name}": value.bfloat16() for name, value in model.experts.named_parameters(prefix="experts") + } + quantized = {f"layers.0.mlp.{name}": value for name, value in export_packed_experts(model).items()} + + merged = Fp8DiffusersExporter().overlay(bf16, quantized) + + for key, value in quantized.items(): + assert torch.equal(merged[key].float(), value.float()) and merged[key].dtype == value.dtype + assert merged["layers.0.mlp.experts.gate_up_proj"].dtype == torch.float8_e4m3fn + assert not any("_quantizer." in key for key in merged) + quantized["layers.0.mlp.experts.down_proj"] = quantized["layers.0.mlp.experts.down_proj"][:, :4] + with pytest.raises(ValueError, match="shape mismatch"): + Fp8DiffusersExporter().overlay(bf16, quantized) + + +@pytest.mark.gpus(1) +def test_wrapper_matches_native_block_on_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for torch._grouped_mm and the Triton routing kernels") + from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock + + torch.manual_seed(0) + block = Qwen3VLMoeTextSparseMoeBlock(_config()) + for param in block.experts.parameters(): + param.data.normal_() + block = block.to(device="cuda", dtype=torch.bfloat16).eval() + native = copy.deepcopy(block) + hidden_states = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16) + token_mask = torch.ones(32, dtype=torch.bool, device="cuda") + token_mask[-5:] = False + + register_packed_experts() + mtq.quantize(block, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda model: model(hidden_states, token_mask=token_mask)) + assert isinstance(block.experts, QuantPackedExperts) + mtq.disable_quantizer(block, "*") + + with torch.no_grad(): + expected, _ = native(hidden_states, token_mask=token_mask) + actual, _ = block(hidden_states, token_mask=token_mask) + assert torch.equal(actual, expected) diff --git a/uv.lock b/uv.lock index 733fd19ce..5b1090b01 100644 --- a/uv.lock +++ b/uv.lock @@ -2388,6 +2388,10 @@ guardrail = [ { name = "retinaface-py" }, { name = "sentencepiece" }, ] +quantization = [ + { name = "datasets" }, + { name = "nvidia-modelopt" }, +] reasoner-eval = [ { name = "attrs" }, { name = "evaluate" }, @@ -2698,6 +2702,7 @@ requires-dist = [ { name = "boto3" }, { name = "botocore", marker = "extra == 'train'" }, { name = "cattrs" }, + { name = "datasets", marker = "extra == 'quantization'" }, { name = "datasets", marker = "extra == 'train'" }, { name = "diffusers", specifier = ">=0.39.0" }, { name = "dists-pytorch", marker = "extra == 'train'" }, @@ -2747,6 +2752,7 @@ requires-dist = [ { name = "nvidia-cudnn-frontend" }, { name = "nvidia-dali-cuda120", marker = "extra == 'train'" }, { name = "nvidia-ml-py" }, + { name = "nvidia-modelopt", marker = "extra == 'quantization'", specifier = "==0.44.0" }, { name = "obstore" }, { name = "omegaconf" }, { name = "open-clip-torch", marker = "extra == 'train'" }, @@ -2812,7 +2818,7 @@ requires-dist = [ { name = "xatlas", marker = "extra == 'train'" }, { name = "zarr", marker = "extra == 'train'" }, ] -provides-extras = ["automl", "automl-prompts", "reasoner-eval", "guardrail", "interactive", "serve", "train"] +provides-extras = ["automl", "automl-prompts", "quantization", "reasoner-eval", "guardrail", "interactive", "serve", "train"] [package.metadata.requires-dev] cu128 = [ @@ -9829,6 +9835,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/72/fb2af0d259a651affdce65fd6a495f0e07a685a0136baf585c5065204ee7/nvidia_ml_py-13.590.48-py3-none-any.whl", hash = "sha256:fd43d30ee9cd0b7940f5f9f9220b68d42722975e3992b6c21d14144c48760e43", size = 50680, upload-time = "2026-01-22T01:14:55.281Z" }, ] +[[package]] +name = "nvidia-modelopt" +version = "0.44.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ninja" }, + { name = "numpy" }, + { name = "nvidia-ml-py" }, + { name = "omegaconf" }, + { name = "packaging" }, + { name = "pulp" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "rich" }, + { name = "safetensors" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "setuptools" }, + { name = "torch", version = "2.10.0", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform == 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.10.0+cu128", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "torch", version = "2.10.0+cu130", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version < '3.13' and sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version < '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "tqdm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/ab/7e12dd238638624cb9d48904e9205abe16a5b26bdd5f9b91e3357821cf90/nvidia_modelopt-0.44.0-py3-none-any.whl", hash = "sha256:9b54a853dfda161db97a0dfce4d7c24269d1d19966b5e2026094186af897f74d", size = 1604658, upload-time = "2026-05-13T20:47:04.007Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.27.5" @@ -11430,6 +11466,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, ] +[[package]] +name = "pulp" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/70/69be07a67621ad804d6cf347965eb4e0d7786a97330d99c31d735aaa6c5a/pulp-3.3.2.tar.gz", hash = "sha256:d0904700c207ac11e25e3b1213b70eae1d6fb25faa719d75f3f15054901258c0", size = 16305346, upload-time = "2026-05-25T09:41:26.207Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/6e/d674f1dde91c71e2ac19e5e5cf1ee6d5845e3aefecd9c39ed9c4b0c9a696/pulp-3.3.2-py3-none-any.whl", hash = "sha256:631b166f72086971a9597f7a0233ababa99bb8d50a01cd543f7758be5a9f86c0", size = 16391742, upload-time = "2026-05-25T09:41:22.2Z" }, +] + [[package]] name = "pure-eval" version = "0.2.3"