Multi-GPU FSDP Support

rapidfireai/automl/automl_utils.py

@@ -7,7 +7,9 @@
 # TODO: add code to validate param_config
 
 
-def get_flattened_config_leaf(param_config: dict[str, Any], prefix: str = "") -> dict[str, Any]:
+def get_flattened_config_leaf(
+    param_config: dict[str, Any], prefix: str = ""
+) -> dict[str, Any]:
     """Flattens the param_config dictionary into a single hierarchy"""
     items = []
     for k, v in param_config.items():
@@ -43,7 +45,9 @@ def get_flattened_config_leaf(param_config: dict[str, Any], prefix: str = "") ->
     return dict(items)
 
 
-def get_runs(param_config: AutoMLAlgorithm | dict[str, Any] | list[Any], seed: int) -> list[dict[str, Any]]:
+def get_runs(
+    param_config: AutoMLAlgorithm | dict[str, Any] | list[Any], seed: int
+) -> list[dict[str, Any]]:
     """Get the runs for the given param_config."""
     # FIXME: how do we handle seed for dict and list?
     if isinstance(param_config, AutoMLAlgorithm):
@@ -56,4 +60,4 @@ def get_runs(param_config: AutoMLAlgorithm | dict[str, Any] | list[Any], seed: i
             config_leaves.extend(get_runs(config, seed))
         return config_leaves
     else:
-        raise ValueError(f"Invalid param_config type: {type(param_config)}")
+        raise ValueError(f"Invalid param_config type: {type(param_config)}")

rapidfireai/automl/base.py

@@ -12,37 +12,47 @@ class AutoMLAlgorithm(ABC):
 
     VALID_TRAINER_TYPES = {"SFT", "DPO", "GRPO"}
 
-    def __init__(self, configs=None, create_model_fn=None, trainer_type: str | None = None, num_runs: int = 1):
+    def __init__(
+        self,
+        configs=None,
+        create_model_fn=None,
+        trainer_type: str | None = None,
+        num_runs: int = 1,
+    ):
         """
         Initialize AutoML algorithm with configurations and trainer type.
-        
+
         Args:
             configs: List of configurations (RFModelConfig for fit mode, dict for evals mode)
             create_model_fn: Optional function to create models (legacy parameter)
             trainer_type: Trainer type ("SFT", "DPO", "GRPO") for fit mode, None for evals mode
             num_runs: Number of runs for random search
-        
+
         Mode detection:
             - If trainer_type is provided: fit mode (requires RFModelConfig instances)
             - If trainer_type is None: evals mode (requires dict instances)
         """
         try:
             self.configs = self._normalize_configs(configs)
             self.num_runs = num_runs
-            
+
             # Detect mode based on trainer_type
             if trainer_type is not None:
                 self.mode = "fit"
                 self.trainer_type = trainer_type.upper()
                 if self.trainer_type not in self.VALID_TRAINER_TYPES:
-                    raise AutoMLException(f"trainer_type must be one of {self.VALID_TRAINER_TYPES}")
+                    raise AutoMLException(
+                        f"trainer_type must be one of {self.VALID_TRAINER_TYPES}"
+                    )
             else:
                 self.mode = "evals"
                 self.trainer_type = None
 
             self._validate_configs()
         except Exception as e:
-            raise AutoMLException(f"Error initializing {self.__class__.__name__}: {e}") from e
+            raise AutoMLException(
+                f"Error initializing {self.__class__.__name__}: {e}"
+            ) from e
 
     def _normalize_configs(self, configs):
         """Normalize configs to list format."""
@@ -56,10 +66,10 @@ def _validate_configs(self):
         """Validate configs based on mode."""
         if not self.configs:
             return
-            
+
         # Import here to avoid circular imports
         from rapidfireai.automl.model_config import RFModelConfig
-        
+
         if self.mode == "fit":
             # Fit mode: must have RFModelConfig instances
             for config in self.configs:

rapidfireai/automl/datatypes.py

@@ -12,12 +12,14 @@ class Range:
 
     def __init__(self, start, end, dtype: str | None = None):
         if dtype is None:
-            self.dtype = "int" if isinstance(start, int) and isinstance(end, int) else "float"
+            self.dtype = (
+                "int" if isinstance(start, int) and isinstance(end, int) else "float"
+            )
         else:
             if dtype not in ("int", "float"):
                 raise ValueError("dtype must be either 'int' or 'float'.")
             self.dtype = dtype
-        if not (isinstance(start, (int, float)) and isinstance(end, (int, float))):
+        if not (isinstance(start, int | float) and isinstance(end, int | float)):
             raise ValueError("start and end must be either int or float.")
         self.start = start
         self.end = end

rapidfireai/automl/grid_search.py

@@ -30,7 +30,7 @@ def recursive_expand_gridsearch(item: Any):
 class RFGridSearch(AutoMLAlgorithm):
     """Grid search algorithm that generates all hyperparameter combinations."""
 
-    def get_runs(self, seed: int=42) -> list[dict[str, Any]]:
+    def get_runs(self, seed: int = 42) -> list[dict[str, Any]]:
         """Generate all possible hyperparameter combinations for grid search."""
         if not isinstance(seed, int) or seed < 0:
             raise AutoMLException("seed must be a non-negative integer")
@@ -58,23 +58,31 @@ def _get_runs_fit(self) -> list[dict[str, Any]]:
 
             for peft_config in peft_configs:
                 peft_instances = (
-                    [{}] if peft_config is None else list(recursive_expand_gridsearch(peft_config._user_params))
+                    [{}]
+                    if peft_config is None
+                    else list(recursive_expand_gridsearch(peft_config._user_params))
                 )
                 training_instances = (
                     [{}]
                     if config.training_args is None
-                    else list(recursive_expand_gridsearch(config.training_args._user_params))
+                    else list(
+                        recursive_expand_gridsearch(config.training_args._user_params)
+                    )
                 )
                 model_kwargs_instances = (
-                    [{}] if config.model_kwargs is None else list(recursive_expand_gridsearch(config.model_kwargs))
+                    [{}]
+                    if config.model_kwargs is None
+                    else list(recursive_expand_gridsearch(config.model_kwargs))
                 )
                 ref_model_kwargs_instances = (
                     [{}]
                     if config.ref_model_kwargs is None
                     else list(recursive_expand_gridsearch(config.ref_model_kwargs))
                 )
                 reward_funcs_instances = (
-                    [{}] if config.reward_funcs is None else list(recursive_expand_gridsearch(config.reward_funcs))
+                    [{}]
+                    if config.reward_funcs is None
+                    else list(recursive_expand_gridsearch(config.reward_funcs))
                 )
 
                 # Get additional kwargs for Trainer
@@ -91,13 +99,18 @@ def _get_runs_fit(self) -> list[dict[str, Any]]:
                     "ref_model_type",
                     "ref_model_kwargs",
                     "reward_funcs",
+                    "num_gpus",
                 }
                 # excluded_attrs = set(config.__dict__.keys()) - set(config.__annotations__.keys())
                 additional_kwargs = {
-                    k: v for k, v in config.__dict__.items() if k not in excluded_attrs and v is not None
+                    k: v
+                    for k, v in config.__dict__.items()
+                    if k not in excluded_attrs and v is not None
                 }
                 additional_kwargs_instances = (
-                    [{}] if not additional_kwargs else list(recursive_expand_gridsearch(additional_kwargs))
+                    [{}]
+                    if not additional_kwargs
+                    else list(recursive_expand_gridsearch(additional_kwargs))
                 )
 
                 # Generate gridsearch combinations
@@ -116,21 +129,26 @@ def _get_runs_fit(self) -> list[dict[str, Any]]:
                                     "model_kwargs": model_kwargs,
                                     "additional_kwargs": additional_kwargs,
                                 }
-
-                                if self.trainer_type == "DPO":
-                                    leaf["ref_model_config"] = {
-                                        "model_name": config.ref_model_name,
-                                        "model_type": config.ref_model_type,
-                                    }
-                                    for ref_model_kwargs in ref_model_kwargs_instances:
-                                        leaf["ref_model_config"]["model_kwargs"] = ref_model_kwargs
-                                        runs.append(leaf)
-                                elif self.trainer_type == "GRPO":
-                                    for reward_func in reward_funcs_instances:
-                                        leaf["reward_funcs"] = reward_func
-                                        runs.append(leaf)
-                                else:
+                                num_gpus = getattr(config, "num_gpus", None)
+                                if num_gpus is not None:
+                                    leaf["num_gpus"] = num_gpus
+
+                            if self.trainer_type == "DPO":
+                                leaf["ref_model_config"] = {
+                                    "model_name": config.ref_model_name,
+                                    "model_type": config.ref_model_type,
+                                }
+                                for ref_model_kwargs in ref_model_kwargs_instances:
+                                    leaf["ref_model_config"][
+                                        "model_kwargs"
+                                    ] = ref_model_kwargs
+                                    runs.append(leaf)
+                            elif self.trainer_type == "GRPO":
+                                for reward_func in reward_funcs_instances:
+                                    leaf["reward_funcs"] = reward_func
                                     runs.append(leaf)
+                            else:
+                                runs.append(leaf)
 
         return runs
 
@@ -147,7 +165,7 @@ def _get_runs_evals(self) -> list[dict[str, Any]]:
                 pipeline = config["pipeline"]
             else:
                 pipeline = None
-                
+
             if pipeline is None:
                 pipelines = [None]
             elif isinstance(pipeline, List):
@@ -156,7 +174,7 @@ def _get_runs_evals(self) -> list[dict[str, Any]]:
                 pipelines = pipeline
             else:
                 pipelines = [pipeline]
-                
+
             for pipeline in pipelines:
                 pipeline_instances = (
                     [{}]
@@ -167,7 +185,10 @@ def _get_runs_evals(self) -> list[dict[str, Any]]:
                 additional_kwargs = {
                     k: v
                     for k, v in config.items()
-                    if k!= "pipeline" and k!= "vllm_config" and k != "openai_config" and v is not None
+                    if k != "pipeline"
+                    and k != "vllm_config"
+                    and k != "openai_config"
+                    and v is not None
                 }
                 additional_kwargs_instances = (
                     [{}]
@@ -181,7 +202,7 @@ def _get_runs_evals(self) -> list[dict[str, Any]]:
                             pipeline_instance = pipeline.__class__(**pipeline_params)
                         else:
                             pipeline_instance = pipeline_params
-                        
+
                         leaf = {
                             "pipeline": pipeline_instance,
                             **additional_kwargs_dict,