Fix cost model num_bytes for sub-byte weight precisions

[dfedoryshchev]

Describe your changes

olive generate-cost-model -p int4 writes a cost model that CaptureSplitInfo cannot read.

The byte count comes from num_params * PRECISON_TO_BYTES[...], and the four sub-byte entries in that table are 0.5, so the product is a float and the column is written as 49250304.0 instead of 49250304. CaptureSplitInfo.split_using_cost_model parses that column with int(), which raises ValueError on the first data row. ValueError is in EXCEPTIONS_TO_RAISE, so the run aborts rather than falling back to another split strategy. int4, uint4, nf4 and fp4 are all argparse choices, and the generate step exits 0 with a success message, so the failure only turns up later in auto-opt.

This rounds the byte count up to an integer. math.ceil rather than int because the value feeds the split_bytes + num_bytes > memory comparison, and a 4-bit layer with an odd parameter count really does occupy the extra half byte. The nine integer-factor precisions are unchanged. The five cost models shipped under olive/assets/cost_models/ are integer in every column already, since they were generated at the fp16 default, which is probably why this has gone unnoticed.

Tests

Adds a regression test that generates a cost model at each of the four sub-byte precisions and asserts the columns parse as integers. It reuses the existing get_cost_model helper and the same tiny model the neighbouring cost-model test uses.

One thing I would rather state than tick: I have not been able to run the suite or lintrunner locally, so the new test is unverified on my side and those two boxes are left unchecked.

Checklist before requesting a review

• Add unit tests for this change.
• Make sure all tests can pass.
• Update documents if necessary. No documentation change needed.
• Lint and apply fixes to your code by running lintrunner -a
• Is this a user-facing change? Yes: generate-cost-model now writes integer byte counts for sub-byte weight precisions, so the generated cost model is usable by CaptureSplitInfo.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟢 Approval recommended

The change is narrowly scoped and addresses a concrete failure mode with added regression coverage; remaining feedback is a minor naming nit.

Pull request overview

This PR fixes generate-cost-model output for sub-byte weight precisions (int4/uint4/nf4/fp4) so the generated CSV can be consumed by the CaptureSplitInfo pass without failing integer parsing, improving reliability of downstream flows like auto-opt.

Changes:

• Round num_bytes up to an integer with math.ceil(...) when generating the cost model CSV (prevents float values like 49250304.0).
• Extend the existing test helper to optionally pass --weight_precision.
• Add a regression test that generates cost models for the four sub-byte precisions and asserts numeric columns parse as integers.

File summaries

File	Description
olive/cli/generate_cost_model.py	Uses math.ceil to ensure integer num_bytes are written to the CSV for sub-byte precisions.
test/passes/pytorch/test_capture_split_info.py	Adds a regression test verifying cost model CSV numeric columns remain integer-parsable for int4/uint4/nf4/fp4.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.