feat: load ModelOpt static-FP8 checkpoints via TorchAO

[pengcuo]

Ports ModelOpt static-FP8 checkpoint loading from imaginaire4 (!10639, merged to i4 main as 39845b6510).

Loads ModelOpt's HF-exported FP8 checkpoints directly: the already-quantized E4M3 weights and their static per-tensor scales are installed as TorchAO tensor subclasses with no dequantization, calibration, or re-quantization. Detection is automatic from the checkpoint's hf_quant_config.json — there is no new flag.

Commits

0ac5e15 — cosmos_framework/utils/generator/quantization.py (+236, pure addition)

• is_modelopt_fp8_checkpoint() — detects quant_method=modelopt / quant_algo=FP8; rejects malformed or unsupported quantization configs rather than silently loading them.
• apply_modelopt_fp8_checkpoint_inplace() — streams the E4M3 weights and their weight_scale / input_scale out of the safetensors shards and builds PrototypeFloat8Tensor weights directly. Follows the root index rather than assuming scale/weight locality (ModelOpt's export can place them in different shards). The full conversion plan is validated — shapes, duplicate targets, missing scales, already-quantized targets — before any module is replaced.
• _ModelOptFloat8Linear — works around two TorchAO 0.16 static-FP8 limits: zero-row inputs (0 // 0 on the block size) and the rank mismatch between a >2D input and its (1, 1) activation scale.

This file is mapped from projects/cosmos3/cosmos3/utils/quantization.py by the release pipeline, so its content is byte-identical to what the next release will emit. It is included because origin/main does not have it yet (last release 2bd01bc predates i4 39845b6510); if a release lands first, this commit can be dropped during rebase.

5bf886b — cosmos_framework/inference/model.py (+87) + tests

_DiffusersLoadPlanner gains defer_modelopt_fp8_weights_loading: the E4M3 weight tensors are dropped from the DCP load plan (and excluded from the missing-key check) so no temporary BF16 copy is ever materialized. from_pretrained_dcp installs them as TorchAO weights once the rest of the checkpoint is in.

Unlike quantization.py, this file is not in the release mapping (see .file_mapping.json) — cosmos_framework/inference/ is maintained here directly, so it is hand-ported.

Divergence from i4 worth reviewing

• i4 factored the load out into a _load_model static method; CF keeps it inline in from_pretrained_dcp (CF has no_dist=, _local_checkpoint_dir threading, and _raise_on_missing_vision_keys that i4 does not). The port follows CF's structure rather than importing i4's refactor.
• CF's from_pretrained_dcp does not apply runtime PTQ — that lives in cosmos_framework/utils/generator/model_loader.py. ModelOpt FP8 checkpoints are therefore supported on the diffusers path only. Wiring the generator path would be a separate change.
• Added a guard i4 does not have: a ModelOpt checkpoint that is not in diffusers layout is rejected with a clear message instead of failing deep inside _diffusers_weight_map with FileNotFoundError.

Behaviour change for non-FP8 checkpoints

is_modelopt_fp8_checkpoint() now runs on every load. A checkpoint carrying an hf_quant_config.json that is not modelopt/FP8 is now rejected with ValueError instead of loading. This is intentional (i4 ef8d577f0a) — an explicit error beats silently loading quantized weights as BF16 — but it is a new failure point.

Everything else on the non-FP8 path is a no-op: the deferral branch is skipped, skipped_target_keys is empty, and the added set lookup never hits.

Guards

ModelOpt FP8 checkpoints are rejected when data_parallel_shard_degree > 1 (weights become tensor subclasses, which requires plain-tensor params, not DTensor shards), when runtime quantization is also requested, and when the layout is not diffusers.

Verification

Run against nvidia/Cosmos3-Experimental@f0cdb8ea cosmos3-nano-fp8-14072026 (21 GB, ModelOpt 0.44.0 export), single H100:

• Loaded 504 calibrated ModelOpt FP8 weights into TorchAO — matches the 504 weight_scale / input_scale pairs in the checkpoint index exactly. The 7 remaining plain nn.Linear are precisely the modules in the checkpoint's ignore list (proj_in, proj_out, time_embedder*, lm_head, visual*).
• Profiler on a converted linear shows aten::_scaled_mm dispatching to nvjet_qqtst_128x128_128x6_1x2_h_bz_TNT, with no aten::mm — real FP8 GEMM, not a dequantize-and-BF16-matmul fallback. Activation quantization is visible as aten::div → aten::clamp → _RoundToFloat8 using the exported input_scale.
• Quantized weights occupy 12.94 GiB vs 25.88 GiB BF16-equivalent.
• inputs/omni/t2i.json generates correctly; after the first compiled step, sampling holds 13.7 it/s (so the _ModelOptFloat8Linear reshape does not break torch.compile).

Non-FP8 regression check on Cosmos3-Nano: loading the same checkpoint on this branch and on main produces an identical SHA256 over all 814 parameters (15.17B elements), and an end-to-end t2i run on this branch is byte-identical to main's output. (Note the sampling pipeline is not reproducible run-to-run, so image hashes alone are not a valid A/B signal — the weight hash is.)

Tests: 18 new (quantization_test.py covers config detection, tensor preservation, ignored/absent targets, and a GPU TorchAO dispatch check; model_test.py covers the planner skip, load/install ordering, and both guards). pyrefly check clean on all four files.

Note: quantization_test.py is new here and is not in the release mapping's include_files. Adding it upstream would be a one-line change to mapping_config.toml; until then a future release adding it would trip the drift gate's collision check once.