Skip to content

feat: wire MXFP8/NVFP4 quantization into inference CLI - #113

Merged
pengcuo merged 2 commits into
mainfrom
liangf/sync-quantization-inference-cli
Jul 15, 2026
Merged

feat: wire MXFP8/NVFP4 quantization into inference CLI#113
pengcuo merged 2 commits into
mainfrom
liangf/sync-quantization-inference-cli

Conversation

@lfengad

@lfengad lfengad commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

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 the model_loader.py hook. But the user-facing OmniInference entrypoint was never wired, so --quantization-method had no effect. This PR connects them.

The missing pieces mapped exactly to MR 10201's packages/cosmos3/cosmos3/* files (→ cosmos_framework/inference/*); the projects/cosmos3/* files were already in.

Changes

  • inference/common/args.py — add QuantizationMethod / QuantizationArgs / QuantizationOverrides; mix them into SetupArgs / SetupOverrides. Fields flow through build_setup's model_dump → model_validate (same as guardrails), so no extra build call is needed.
  • inference/model.py — add Cosmos3OmniConfig.quantization property + setter and thread quantization_config through from_pretrained_dcp. Uses the local direct-construction style unstructure_config(QuantizationConfig(**value)) rather than upstream's LazyCall wrapper, matching the sibling parallelism/compile setters in this repo.
  • inference/inference.py — add _get_quantization_config and pass it into both _create branches (load_model_from_checkpoint and from_pretrained_dcp).

Effect

  • Standard experiment DCP path (load_model_from_checkpoint): quantization applied end-to-end via apply_quantization_inplace. ✅
  • HF from_pretrained_dcp path: stores config.quantization without an apply hook — matches upstream MR behavior (not a regression introduced here).
  • Only supported on Blackwell + replicated inference (no FSDP sharding); guarded by the existing model_loader.py DP-shard check.

Not included

  • websocket_policy_server/serve_policy.py + README from the MR — that action-eval subtree is not present in this repo.
  • The MR's cosmetic blank-line change in packages/cosmos3/cosmos3/args.py.

Verification

  • py_compile passes on all three files.
  • Load-bearing pydantic logic verified in-container (isolated replica): defaults, mutable-list-default isolation, build_quantization, Literal rejection of invalid methods, MRO/inheritance, and SetupOverrides → SetupArgs field flow — all assertions pass.
  • Full module import not exercised: cosmos-framework requires Python ≥3.12 while the available i4 container env is 3.10; needs the cosmos-framework 3.13 venv to import the real module.

🤖 Generated with Claude Code

lfengad and others added 2 commits July 14, 2026 21:23
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>
@pengcuo
pengcuo merged commit ace6140 into main Jul 15, 2026
8 checks passed
@pengcuo
pengcuo deleted the liangf/sync-quantization-inference-cli branch July 15, 2026 06:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants