Merge pull request #6 from pytorch-labs/fmassa/fix_compute_cost

fmassa · web-flow · commit b51cd9a24f42 · 2025-06-13T18:43:58.000+02:00
Fix computation cost to take number of arguments into account
diff --git a/autoparallel/optimize_sharding.py b/autoparallel/optimize_sharding.py
@@ -100,7 +100,7 @@ def build_ds(self):
                         )
                         ds[(s_i, argi, ss, ii)] = {
                             "va": va,
-                            "cost": comm_cost + compute_cost,
+                            "cost": comm_cost + compute_cost / num_args[s_i],
                             "full_strat": ssi,
                             "out_strat": ssi.output_specs,
                             "inp_strat": ssi.input_specs[argi],

Original file line number	Diff line number	Diff line change
`@@ -100,7 +100,7 @@ def build_ds(self):`
`100`	`100`	`)`
`101`	`101`	`ds[(s_i, argi, ss, ii)] = {`
`102`	`102`	`"va": va,`
`103`		`- "cost": comm_cost + compute_cost,`
	`103`	`+ "cost": comm_cost + compute_cost / num_args[s_i],`
`104`	`104`	`"full_strat": ssi,`
`105`	`105`	`"out_strat": ssi.output_specs,`
`106`	`106`	`"inp_strat": ssi.input_specs[argi],`