Fix all replicas had the same _llm_identifier for `CudaDevicePlacem…

…entMixin` (#941) * Fix CUDA device placement with multiple replicas * Print replica id * Copy `step` for each replica
argilla-io · Sep 2, 2024 · 56b4036 · 56b4036
1 parent c8f4d61
commit 56b4036
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/src/distilabel/llms/mixins/cuda_device_placement.py b/src/distilabel/llms/mixins/cuda_device_placement.py
@@ -207,11 +207,6 @@ def _get_cuda_device(self, device_map: Dict[str, List[int]]) -> Union[int, None]
                 return device
 
         return None
-        raise RuntimeError(
-            "Couldn't find an available CUDA device automatically to be used by the LLM"
-            f" '{self._llm_identifier}'. For forcing the use of a specific device, set the"
-            " `cuda_devices` attribute to a list with the desired device(s)."
-        )
 
     def _set_cuda_visible_devices(self) -> None:
         """Sets the `CUDA_VISIBLE_DEVICES` environment variable to the list of CUDA devices

diff --git a/src/distilabel/pipeline/base.py b/src/distilabel/pipeline/base.py
@@ -1222,8 +1222,14 @@ def _run_steps(self, steps: Iterable[str]) -> None:
 
             step_num_replicas: int = step.resources.replicas if step.is_normal else 1  # type: ignore
             for replica in range(step_num_replicas):
-                self._logger.debug(f"Running 1 replica of step '{step.name}'...")
-                self._run_step(step=step, input_queue=input_queue, replica=replica)
+                self._logger.debug(
+                    f"Running 1 replica of step '{step.name}' with ID {replica}..."
+                )
+                self._run_step(
+                    step=step.model_copy(deep=True),
+                    input_queue=input_queue,
+                    replica=replica,
+                )
 
     def _add_batches_back_to_batch_manager(self) -> None:
         """Add the `Batch`es that were sent to a `Step` back to the `_BatchManager`. This

diff --git a/src/distilabel/pipeline/step_wrapper.py b/src/distilabel/pipeline/step_wrapper.py
@@ -76,7 +76,7 @@ def _init_cuda_device_placement_mixin(attr: CudaDevicePlacementMixin) -> None:
                 attr.disable_cuda_device_placement = True
             else:
                 desired_num_gpus = self.step.resources.gpus or 1
-                attr._llm_identifier = self.step.name
+                attr._llm_identifier = f"{self.step.name}-replica-{self.replica}"
                 attr._desired_num_gpus = desired_num_gpus
 
         for field_name in self.step.model_fields_set: