fix: main eval loop zips skipped combos and omits calibrated kwargs

[andrewwhitecdw]

This PR addresses the following issue in fouroversix/scripts/ptq/coordinators/modal.py: main eval loop zips skipped combos and omits calibrated kwargs.

Changes

• fouroversix/scripts/ptq/coordinators/modal.py: main eval loop zips skipped combos and omits calibrated kwargs.

Details

--- a/fouroversix/scripts/ptq/coordinators/modal.py
+++ b/fouroversix/scripts/ptq/coordinators/modal.py
@@ -1,38 +1,46 @@
-        models_and_ptq_methods = list(itertools.product(model_names, ptq_methods))
-        function_calls = []
-
-        for model_name, ptq_method in models_and_ptq_methods:
-            tasks_to_evaluate = self.get_tasks_to_evaluate(
-                model_name,
-                ptq_method,
-                tasks,
-            )
-
-            if len(tasks_to_evaluate) == 0:
-                continue
-
-            evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu)
-
-            calibrated_kwargs = evaluator_cls.get_calibrated_kwargs(
-                model_name,
-                self.get_session(),
-                **kwargs,
-            )
-
-            function_calls.append(
-                evaluator_cls().evaluate_on_modal.spawn(
-                    model_name=model_name,
-                    tasks=tasks_to_evaluate,
-                    save_path=FOUROVERSIX_CACHE_PATH / "ptq",
-                    **{**kwargs, **calibrated_kwargs},
-                ),
-            )
-
-        all_results = modal.FunctionCall.gather(*function_calls)
-
-        for (model_name, ptq_method), results in zip(
-            models_and_ptq_methods,
-            all_results,
-            strict=True,
-        ):
-            self.save_results(model_name, ptq_method, kwargs, results)
+        models_and_ptq_methods = list(itertools.product(model_names, ptq_methods))
+        evaluated_models_and_ptq_methods = []
+        evaluated_kwargs = []
+        function_calls = []
+
+        for model_name, ptq_method in models_and_ptq_methods:
+            tasks_to_evaluate = self.get_tasks_to_evaluate(
+                model_name,
+                ptq_method,
+                tasks,
+            )
+
+            if len(tasks_to_evaluate) == 0:
+                continue
+
+            evaluated_models_and_ptq_methods.append((model_name, ptq_method))
+
+            evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu)
+
+            calibrated_kwargs = evaluator_cls.get_calibrated_kwargs(
+                model_name,
+                self.get_session(),
+                **kwargs,
+            )
+
+            merged_kwargs = {**kwargs, **calibrated_kwargs}
+            evaluated_kwargs.append(merged_kwargs)
+
+            function_calls.append(
+                evaluator_cls().evaluate_on_modal.spawn(
+                    model_name=model_name,
+                    tasks=tasks_to_evaluate,
+                    save_path=FOUROVERSIX_CACHE_PATH / "ptq",
+                    **merged_kwargs,
+                ),
+            )
+
+        all_results = modal.FunctionCall.gather(*function_calls)
+
+        for (model_name, ptq_method), merged_kwargs, results in zip(
+            evaluated_models_and_ptq_methods,
+            evaluated_kwargs,
+            all_results,
+            strict=True,
+        ):
+            self.save_results(model_name, ptq_method, merged_kwargs, results)

Tests

Let me know if you want tests added for this fix or not.