From 6f2ff3ea0c4c2c5f24a1ef61ac35d015ede8d5fa Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 14 Jul 2026 21:23:30 -0700 Subject: [PATCH 1/2] feat: wire MXFP8/NVFP4 quantization into inference CLI Ports the inference-CLI plumbing from imaginaire4 MR 10201 that was missing on this side. The quantization backend (QuantizationConfig, apply_quantization_inplace, and the model_loader hook) was already synced; this connects it to the user-facing OmniInference entrypoint so --quantization-method actually takes effect. - inference/common/args.py: add QuantizationMethod / QuantizationArgs / QuantizationOverrides; mix them into SetupArgs / SetupOverrides so the fields flow through build_setup's model_dump -> model_validate. - inference/model.py: add Cosmos3OmniConfig.quantization property and thread quantization_config through from_pretrained_dcp. Uses the local direct-construction style (unstructure_config(QuantizationConfig(**v))) rather than upstream's LazyCall wrapper, matching the sibling parallelism/compile setters here. - inference/inference.py: add _get_quantization_config and pass it into both _create branches (load_model_from_checkpoint and from_pretrained_dcp). Effect: the standard experiment DCP path applies quantization end-to-end (load_model_from_checkpoint -> apply_quantization_inplace). The HF from_pretrained_dcp path stores the config without an apply hook, matching upstream. The websocket_policy_server change from the MR is intentionally skipped (that action-eval subtree is not in this repo). Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/common/args.py | 36 +++++++++++++++++++++-- cosmos_framework/inference/inference.py | 12 ++++++++ cosmos_framework/inference/model.py | 15 ++++++++++ 3 files changed, 61 insertions(+), 2 deletions(-) diff --git a/cosmos_framework/inference/common/args.py b/cosmos_framework/inference/common/args.py index 4fa87c97..13ee789e 100644 --- a/cosmos_framework/inference/common/args.py +++ b/cosmos_framework/inference/common/args.py @@ -603,6 +603,38 @@ def build_checkpoint(self, *, checkpoints: dict[str, CheckpointConfig]) -> Check CfgpSize = Annotated[int, pydantic.Field(ge=1, le=2)] CompiledRegion = Literal["all", "language"] +# Low-precision quantization method to apply to the model at load time. +# One of ``mxfp8`` / ``nvfp4``, or ``None`` (default) to disable. +# Routed to the VFM model loader, which selects an FSDP-compatible +# (module-swap) path when sharded (``dp_shard_size > 1``) and an in-place +# path when replicated (``dp_shard_size == 1``). Note ``mxfp8`` / ``nvfp4`` +# are only supported on the replicated path. +QuantizationMethod = Literal["mxfp8", "nvfp4"] + + +class QuantizationArgs(ArgsBase): + """Low-precision quantization arguments applied to the model at load time.""" + + quantization_method: QuantizationMethod | None + quantization_include_regex: list[str] + quantization_exclude_regex: list[str] + + +class QuantizationOverrides(OverridesBase): + quantization_method: QuantizationMethod | None = None + """Quantization method (``mxfp8`` / ``nvfp4``), or ``None`` to disable. + + Post-training quantization (PTQ) is applied in-place to the model at load + time. Only supported on Blackwell architectures and when FSDP sharding is disabled. + """ + quantization_include_regex: list[str] = ["language_model.model.layers"] + """Regexes matched against module FQNs; a Linear is quantized only if it matches one (empty = all).""" + quantization_exclude_regex: list[str] = pydantic.Field(default_factory=list) + """Regexes matched against module FQNs; a Linear is skipped if it matches any.""" + + def build_quantization(self) -> QuantizationArgs: + return self._build(QuantizationArgs) + class ParallelismArgs(ArgsBase): """Parallelism arguments.""" @@ -702,7 +734,7 @@ class GuardrailOverrides(OverridesBase): """Offload guardrail models to CPU.""" -class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, GuardrailArgs): +class SetupArgs(ABC, CheckpointArgs, ParallelismArgs, QuantizationArgs, GuardrailArgs): output_dir: ResolvedPath keep_going: bool skip_invalid_samples: bool @@ -737,7 +769,7 @@ def get_variant(cls) -> str: return cls.model_fields["variant"].default -class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, GuardrailOverrides): +class SetupOverrides(ABC, CheckpointOverrides, ParallelismOverrides, QuantizationOverrides, GuardrailOverrides): """Inference setup arguments.""" output_dir: Annotated[ResolvedPath | None, tyro.conf.arg(aliases=("-o",))] = None diff --git a/cosmos_framework/inference/inference.py b/cosmos_framework/inference/inference.py index a7f338ca..30bd04b5 100644 --- a/cosmos_framework/inference/inference.py +++ b/cosmos_framework/inference/inference.py @@ -23,6 +23,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig from cosmos_framework.inference.args import ( ModelMode, NegativeMetadataMode, @@ -1055,6 +1056,14 @@ def _get_compile_config(cls, setup_args: ParallelismArgs) -> CompileConfig: compile_dynamic=setup_args.compile_dynamic, ) + @classmethod + def _get_quantization_config(cls, setup_args: SetupArgs) -> QuantizationConfig: + return QuantizationConfig( + method=setup_args.quantization_method, + include_regex=list(setup_args.quantization_include_regex), + exclude_regex=list(setup_args.quantization_exclude_regex), + ) + @override @classmethod def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: @@ -1064,6 +1073,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: sampler_override = setup_args.sampler parallelism_config = cls._get_parallelism_config(setup_args) compile_config = cls._get_compile_config(setup_args) + quantization_config = cls._get_quantization_config(setup_args) if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE: from cosmos_framework.inference.common.config import save_config from cosmos_framework.utils.generator.model_loader import load_model_from_checkpoint @@ -1081,6 +1091,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: credential_path=setup_args.credential_path or None, parallelism_config=attrs.asdict(parallelism_config), compile_config=attrs.asdict(compile_config), + quantization_config=attrs.asdict(quantization_config), load_ema_to_reg=setup_args.use_ema_weights, experiment_opts=[ *setup_args.experiment_overrides, @@ -1130,6 +1141,7 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: config=config, parallelism_config=parallelism_config, compile_config=compile_config, + quantization_config=quantization_config, ).model if model.config.rectified_flow_inference_config.scheduler_type != sampler_override: model.config.rectified_flow_inference_config.scheduler_type = sampler_override diff --git a/cosmos_framework/inference/model.py b/cosmos_framework/inference/model.py index e9c38b48..042bab5b 100644 --- a/cosmos_framework/inference/model.py +++ b/cosmos_framework/inference/model.py @@ -34,6 +34,7 @@ from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig +from cosmos_framework.configs.base.defaults.quantization import QuantizationConfig from cosmos_framework.inference.common.args import CheckpointType from cosmos_framework.inference.common.checkpoints import register_checkpoints from cosmos_framework.inference.common.config import structure_config, undo_config_dict_replacements, unstructure_config @@ -416,6 +417,16 @@ def compile(self, value: dict | None): return self.model.setdefault("config", {})["compile"] = unstructure_config(CompileConfig(**value)) + @property + def quantization(self) -> dict: + return self.model.get("config", {}).get("quantization", {}) + + @quantization.setter + def quantization(self, value: dict | None): + if value is None: + return + self.model.setdefault("config", {})["quantization"] = unstructure_config(QuantizationConfig(**value)) + class Cosmos3OmniModel(transformers.PreTrainedModel): config_class = Cosmos3OmniConfig # type: ignore @@ -448,6 +459,7 @@ def from_pretrained_dcp( config: Cosmos3OmniConfig | None = None, parallelism_config: ParallelismConfig | None = None, compile_config: CompileConfig | None = None, + quantization_config: QuantizationConfig | None = None, ): if config is None: config = Cosmos3OmniConfig.from_pretrained(checkpoint_path) @@ -455,8 +467,11 @@ def from_pretrained_dcp( parallelism_config = ParallelismConfig() if compile_config is None: compile_config = CompileConfig() + if quantization_config is None: + quantization_config = QuantizationConfig() config.parallelism = attrs.asdict(parallelism_config) config.compile = attrs.asdict(compile_config) + config.quantization = attrs.asdict(quantization_config) model = cls(config) checkpoint_type = CheckpointType.from_path(checkpoint_path) match checkpoint_type: From 86bede74e5592a58cccecd5e8f6ea4bf7d0b0245 Mon Sep 17 00:00:00 2001 From: "liang.feng" Date: Tue, 14 Jul 2026 22:01:17 -0700 Subject: [PATCH 2/2] fix: register public alias for QuantizationConfig from_pretrained_dcp now always sets config.quantization (default QuantizationConfig), so the serialized model config carries a quantization block. build_public_model_config walks the config and maps each sub-config's type path to a public alias; without an entry for QuantizationConfig it raised "No public alias registered for type path", breaking convert_model_to_dcp / export_model (and every training job that runs them). Register the alias mirroring parallelism_config / compile_config. This is a cosmos-framework OSS-remap-layer registration, not present in the upstream i4 MR. Round-trip (build -> restore) verified. Co-Authored-By: Claude Opus 4.8 (1M context) --- cosmos_framework/inference/common/public_model_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/cosmos_framework/inference/common/public_model_config.py b/cosmos_framework/inference/common/public_model_config.py index 584a65c9..aff1a057 100644 --- a/cosmos_framework/inference/common/public_model_config.py +++ b/cosmos_framework/inference/common/public_model_config.py @@ -37,6 +37,7 @@ "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowInferenceConfig": "rectified_flow_inference_config", "projects.cosmos3.vfm.configs.base.defaults.model_config.RectifiedFlowTrainingConfig": "rectified_flow_training_config", "projects.cosmos3.vfm.configs.base.defaults.parallelism.ParallelismConfig": "parallelism_config", + "projects.cosmos3.vfm.configs.base.defaults.quantization.QuantizationConfig": "quantization_config", "projects.cosmos3.vfm.configs.base.defaults.vlm.PretrainedWeightsConfig": "pretrained_weights_config", "projects.cosmos3.vfm.configs.base.defaults.vlm.VLMConfig": "vlm_config", }