feat: wire MXFP8/NVFP4 quantization into inference CLI - #113
Merged
Conversation
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
foreverlms
approved these changes
Jul 15, 2026
pengcuo
approved these changes
Jul 15, 2026
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
Ports the inference-CLI plumbing from imaginaire4 MR 10201 (
Inference optimizations [4/n]: MXFP8 / NVFP4 PTQ) that was missing on this side.The quantization backend was already synced (via the 2026-07-10 release):
QuantizationConfig(configs/base/defaults/quantization.py),apply_quantization_inplace(utils/generator/quantization.py), and themodel_loader.pyhook. But the user-facingOmniInferenceentrypoint was never wired, so--quantization-methodhad no effect. This PR connects them.The missing pieces mapped exactly to MR 10201's
packages/cosmos3/cosmos3/*files (→cosmos_framework/inference/*); theprojects/cosmos3/*files were already in.Changes
inference/common/args.py— addQuantizationMethod/QuantizationArgs/QuantizationOverrides; mix them intoSetupArgs/SetupOverrides. Fields flow throughbuild_setup'smodel_dump → model_validate(same as guardrails), so no extra build call is needed.inference/model.py— addCosmos3OmniConfig.quantizationproperty + setter and threadquantization_configthroughfrom_pretrained_dcp. Uses the local direct-construction styleunstructure_config(QuantizationConfig(**value))rather than upstream'sLazyCallwrapper, matching the siblingparallelism/compilesetters in this repo.inference/inference.py— add_get_quantization_configand pass it into both_createbranches (load_model_from_checkpointandfrom_pretrained_dcp).Effect
load_model_from_checkpoint): quantization applied end-to-end viaapply_quantization_inplace. ✅from_pretrained_dcppath: storesconfig.quantizationwithout an apply hook — matches upstream MR behavior (not a regression introduced here).model_loader.pyDP-shard check.Not included
websocket_policy_server/serve_policy.py+ README from the MR — that action-eval subtree is not present in this repo.packages/cosmos3/cosmos3/args.py.Verification
py_compilepasses on all three files.build_quantization,Literalrejection of invalid methods, MRO/inheritance, andSetupOverrides → SetupArgsfield flow — all assertions pass.🤖 Generated with Claude Code