diff --git a/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py b/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py index e812b95..339853f 100644 --- a/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py +++ b/runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py @@ -508,6 +508,28 @@ def reweight( for batch_size, count in histogram.items() if int(batch_size) in by_batch } + covered = sum(supported.values()) + if covered == 0: + return {"coverage": 0.0, "components_ms": None, "unsupported": histogram} + values = { + category: sum( + by_batch[batch_size]["components_ms"][category] * count + for batch_size, count in supported.items() + ) + / covered + for category in (*CATEGORIES, "total") + } + return { + "coverage": covered / total_samples, + "supported_samples": covered, + "total_samples": total_samples, + "components_ms": values, + "unsupported": { + batch_size: count + for batch_size, count in histogram.items() + if int(batch_size) not in by_batch + }, + } def reweight_joint( @@ -544,30 +566,6 @@ def reweight_joint( if state not in by_state }, } - covered = sum(supported.values()) - if covered == 0: - return {"coverage": 0.0, "components_ms": None, "unsupported": histogram} - values = { - category: sum( - by_batch[batch_size]["components_ms"][category] * count - for batch_size, count in supported.items() - ) - / covered - for category in (*CATEGORIES, "total") - } - return { - "coverage": covered / total_samples, - "supported_samples": covered, - "total_samples": total_samples, - "components_ms": values, - "unsupported": { - batch_size: count - for batch_size, count in histogram.items() - if int(batch_size) not in by_batch - }, - } - - def subtract(right: dict[str, float], left: dict[str, float]) -> dict[str, float]: return {name: right[name] - left[name] for name in left} diff --git a/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py b/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py index b75e17d..e13dba0 100644 --- a/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py +++ b/runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py @@ -250,6 +250,12 @@ class FidelityEnvelopeTest(unittest.TestCase): self.assertEqual(module.op_category(name), category) with self.assertRaisesRegex(ValueError, "unclassified"): module.op_category("unknown_graph_overhead") + component_row = {name: 1.0 for name in (*module.CATEGORIES, "total")} + reweighted = module.reweight( + {4: {"components_ms": component_row}}, {"4": 3} + ) + self.assertEqual(reweighted["coverage"], 1.0) + self.assertEqual(reweighted["components_ms"]["total"], 1.0) def test_materialize_qwen235_allreduce_requires_serving_contract(self) -> None: module = load("materialize_qwen235_v020_allreduce.py")