LLM-VM does not support multiple GPUs currently

[MehmetMHY]

Currently LLM-VM does not support multiple GPU setups. Using runpod, I rented a setup with 2 RTX 3090 GPUs. Well running the local Bloom model example from the docs. I ran into this error:

`EleutherAI/pythia-70m-deduped` loaded on 2 GPUs.
Using model: bloom
Running with an empty context
Exception in thread Thread-2 (new_thread):
Traceback (most recent call last):
  File "/usr/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.10/dist-packages/llm_vm/completion/optimize.py", line 45, in new_thread
    t[0] = foo()
  File "/usr/local/lib/python3.10/dist-packages/llm_vm/completion/optimize.py", line 259, in promiseCompletion
    best_completion = self.call_big(prompt, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/llm_vm/client.py", line 102, in CALL_BIG
    return self.teacher.generate(prompt, max_len,**kwargs)
  File "/usr/local/lib/python3.10/dist-packages/llm_vm/onsite_llm.py", line 153, in generate
    generate_ids=self.model.generate(inputs.input_ids, max_length=max_length, **generation_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1695, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'DataParallel' object has no attribute 'generate'
{'status': 0, 'resp': 'cannot unpack non-iterable NoneType object'}

On fix I had for this error was by adding the following code changes to src/llm_vm/onsite_llm.py (added the if-else statement under the "account ofr cases where model is wrapped..." comment):

diff --git a/src/llm_vm/onsite_llm.py b/src/llm_vm/onsite_llm.py
index 9fcfe3c..613acbf 100644
--- a/src/llm_vm/onsite_llm.py
+++ b/src/llm_vm/onsite_llm.py

@@ -151,7 +141,13 @@ class BaseOnsiteLLM(ABC):
             inputs = self.tokenizer(prompt, return_tensors="pt", **tokenizer_kwargs).to(device[0])
         else:
             inputs = self.tokenizer(prompt, return_tensors="pt").to(device)
-        generate_ids=self.model.generate(inputs.input_ids, max_length=max_length, **generation_kwargs)
+        
+        # account for cases where the model is wrapped in DataParallel
+        if isinstance(self.model, torch.nn.DataParallel):
+            generate_ids = self.model.module.generate(inputs.input_ids, max_length=max_length, **generation_kwargs)
+        else:
+            generate_ids = self.model.generate(inputs.input_ids, max_length=max_length, **generation_kwargs)
+
         resp= self.tokenizer.batch_decode(generate_ids,skip_special_tokens=True,clean_up_tokenization_spaces=False)[0]
         # need to drop the len(prompt) prefix with these sequences generally
         # because they include the prompt.

☝️ This change resolves the error but it does not fix the core issue. In this change we are using DataParallel.module to access the model's generate() function. In doing so, we would be bypassing the DataParallel wrapper and skipping the parallelism provided by "DataParallel". I believe this is the case and that we should implement a new solution.

[VictorOdede]

This issue can be solved using Ray.io. For distributed training: https://docs.ray.io/en/latest/train/train.html For distributed inference: https://docs.ray.io/en/latest/serve/index.html

[VictorOdede]

Also look into: https://huggingface.co/docs/accelerate/index

[Aryan8912]

can you please assign to me

[berkay3500]

Is this issue still open?