[Bounty] PyTorch & HuggingFace Interface #139

risingsunomi · 2024-08-11T02:24:36Z

Hello all,

I’ve made some updates to the exo library based on the bounty mentioned in this tweet/X post. These changes aim to integrate PyTorch and expand access to various language models through Hugging Face’s AutoModelForCausalLM.

What's New?

ShardedHuggingFaceModel: Adds sharding support for Hugging Face models.
PyTorchDynamicShardInferenceEngine: A new inference engine that uses PyTorch tensors for dynamic sharding.

These updates enable the exo library to use PyTorch, allowing access to a broader range of language models.

Limitations and Bugs

Right now the ShardedHuggingFaceModel is focused on using LlamaForCausalLM from the huggingface transformers library. From that model we break it up using LLamaModel and the layers it contains. We can then select the layers and run the pytorch tensors over them as need. I focused on using llama3.1 8B as I could only slightly run that.

Due to my current hardware limitations (specifically GPU and VRAM), I wasn’t able to fully test this across multiple nodes. The model currently takes about 30 seconds per token to generate for me (I have slow GPUs), which might be related to the absence of caching (not implemented due to VRAM constraints). It’s running without reaching an EOT and the outputs seem random.

Request for Feedback

I’m sharing this in the hope that others can test it on more capable setups and provide feedback on how to enhance performance and stability.

Important Note on Meta LLaMA 3.1 Model

If you plan to test with the official Meta LLaMA 3.1 model, please note:

Access: You’ll need to request access and authenticate using huggingface-cli to download it.
Command: Run the following command before using the model:
```
huggingface-cli login
```
I’m exploring ways to simplify this process, but for now, it’s necessary.

Chat API Update

Added an option to select the LLaMA 3.1 model in the chat API.

Looking forward to any feedback or suggestions you might have.

Thank you

AlexCheema · 2024-08-18T19:09:23Z

Hey, sorry for the delay. I haven't had a chance to check this properly yet. I'll be able to look next week.

risingsunomi · 2024-08-18T22:35:31Z

Hey, sorry for the delay. I haven't had a chance to check this properly yet. I'll be able to look next week.

Sounds good. Let me know anything needed. Thank you

AlexCheema · 2024-08-23T17:08:49Z

exo/inference/pytorch/model/hf.py

+
+        # self.past_key_values = DynamicCache()
+
+    def forward_layers(


I really like this approach of generalising this so it works for other models without having to explicitly implement them.

Can you write a test for a model with a different architecture to make sure this generalises e.g. recurrent Gemma?
I wonder if we need a little bit of model-specific behaviour to enable this in general?

AlexCheema · 2024-08-23T17:09:17Z

exo/inference/pytorch/helpers.py

+            async for chunk in response.content.iter_chunked(8192):
+                f.write(chunk)
+
+async def download_files(urls: List[str], output_paths: List[Path]):


No, I can remove

Please remove

Please remove

Sorry forgot to. Will do that now.

AlexCheema · 2024-08-23T17:11:26Z

exo/inference/pytorch/inference.py

+        self.shard = None
+        self.model = None
+        self.tokenizer = None
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


Are these the only options? I think supporting e.g. Mac with mps would be great since then you can run heterogeneous clusters.

One thing to try at some point would be mixing MLX and PyTorch and see if they are interoperable with exactly the same model.

With pytorch I don't think mac is fully rolled out yet. There seems to be some work arounds but CUDA and CPU are the only options on the pytorch download website. pytorch even stopped ROCm support for AMD

They have a nightly for testing MPS https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/

What about this in the official "stable" docs: https://pytorch.org/docs/stable/notes/mps.html

Will try that but currently no mac to test. When I get through these other fixes though I can definitely add it for you or other mac users to test.

AlexCheema · 2024-08-23T17:12:22Z

exo/inference/pytorch/model/hf.py

+        # Load the model
+        self.full_model = AutoModelForCausalLM.from_pretrained(
+            shard.model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,


Again, are these the only options? Would want support across other platforms

AlexCheema · 2024-08-23T17:13:27Z

exo/inference/pytorch/model/hf.py

+
+            layers.append(layer)
+
+        self.full_model.model.layers = nn.ModuleList(layers)


What does the peak memory usage look like here? I'm not sure of the specifics of python if this is going to hold each layer twice. Not sure but perhaps setting them in place would be more memory efficient.

They shouldn't be held twice as when the ensure_shard function is called in the infer_prompt or infer_tensor the init class function is called which loads the needed layers each time depending on the shard. Will make sure about memory limits though and usage.

AlexCheema · 2024-08-23T17:15:12Z

exo/inference/pytorch/model/hf.py

+
+
+        # Load the model
+        self.full_model = AutoModelForCausalLM.from_pretrained(


Will this download the entire model?
We have code to selectively download the model from HuggingFace so you don't have to download all layers on every device: exo/download

Also this won't work with our download progress code. We show in the TUI what the download progress of the model is.

Will this download the entire model? We have code to selectively download the model from HuggingFace so you don't have to download all layers on every device: exo/download

Will look at using that code because yes it currently does download all the model

AlexCheema · 2024-08-23T17:16:01Z

exo/inference/pytorch/model/utils.py

+import torch
+from torch.nn import functional as F
+
+def sample_logits(logits, temperature=1.0, top_k=0, top_p=1.0, alpha_f=0.0, alpha_p=0.0):


Can this be imported from somewhere rather than copy-pasta into the codebase? It looks like boilerplate code from somewhere.

Yes, I was testing it as the default values but will clean that part up. I will set it in the Interface class settings to be used.

AlexCheema · 2024-08-23T17:16:31Z

install.ps1

+& .\.venv\Scripts\Activate.ps1
+
+# Install the package in the virtual environment
+pip install .


Windows? Did this work on windows? Curious if it works there.

I was testing on Windows but couldn't fully get it working right. Will test again and make sure as I switched to using Linux to do further dev

@AlexCheema I can confirm that it is working on Windows, but there are a few issues:

PyTorch 2.4 doesn't install. In order to get it working, it needs the nightly build

In main.py

for s in [signal.SIGINT, signal.SIGTERM]: loop.add_signal_handler(s, handle_exit)

isn't supported on windows.

If I change handle_exit() to:

def handle_exit(): asyncio.ensure_future(shutdown(loop)) if platform.system() != "Windows": for s in [signal.SIGINT, signal.SIGTERM]: loop.add_signal_handler(s, handle_exit) else: # On Windows, we can only reliably catch SIGINT (Ctrl+C) signal.signal(signal.SIGINT, lambda signum, frame: handle_exit())

Getting some kind of network error between the GUI and the backend

It seems to work. It's insanely slow for me though (no GPU... the raspberry pis are much faster 😄). Windows changes perhaps out of scope for this PR though.

Wowzah!

This is exciting. A lot of the community were mad we didn't support Windows. We have a bounty here if you want to get it working there once this one is merged: #186

AlexCheema · 2024-08-23T17:18:40Z

Great work. You clearly thought about this and implemented a really nice solution. I particularly like the generalisation of model splitting, rather than doing each one separately.

Take a look through the comments I left.

AlexCheema · 2024-08-23T17:19:59Z

The main thing I want to address and test is device support. We can make this the default inference engine if it works reliably across many devices.

On that point, if we can automate the bootstrapping of the environment for each user (e.g. install drivers, whatever else is needed to run on their device) that would be great. We don't have to do this in this PR/bounty, we can do another. But I would love to discuss and figure out how this can best be done.

AlexCheema · 2024-08-23T17:21:45Z

exo/inference/pytorch/test_build_transformer.py

@@ -0,0 +1,21 @@
+import unittest


Run this test in circle ci ./.circleci/config.yml

AlexCheema · 2024-08-23T17:21:49Z

exo/inference/pytorch/test_inference_engine.py

@@ -0,0 +1,33 @@
+


Run this test in circle ci ./.circleci/config.yml

AlexCheema · 2024-08-23T17:22:33Z

exo/inference/pytorch/test_inference_engine.py

+        n_layers=12
+    )
+
+    engine = PyTorchDynamicShardInferenceEngine(


Is this test complete? We need a test that tests the model splitting. Take a look at exo/inference/test_inference_engine.py. You can just add the test there.

risingsunomi

Reviewed and looks good. Will work through notes to improve

…unless wanting to rebuild the whole model

…s, update infer to pass logits and hidden_states

…g, added model value options

…o it will stop, context was also dropping due to not having all the logits

…into pr139-dev-oct24

…nodes

…del layers as that will not be passed to other nodes

…en and not max_new_tokens

…tion, working on exo inference engine

…ting split model test, trying to get inference engine to hit stop token

…dev-oct24

Pr139 dev oct24

risingsunomi · 2024-12-01T01:51:18Z

I have done some updates to the engine but doing some more testing now. I am running into some breaks and trying to run a larger model. Will notify when more testable.

…rt for more models like Qwen as building Qwen engine, updated full llama test with more complex prompts

risingsunomi · 2024-12-01T10:58:54Z

Should be good to test now. Going to try on larger models. In the future, will write the qwen engine pending this one is all good.

AlexCheema reviewed Aug 23, 2024

View reviewed changes

exo/inference/pytorch/test_inference_engine.py Outdated

@@ -0,0 +1,33 @@

Copy link

Contributor

AlexCheema Aug 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Run this test in circle ci ./.circleci/config.yml

AlexCheema reviewed Aug 23, 2024

View reviewed changes

risingsunomi commented Aug 23, 2024

View reviewed changes

This was referenced Aug 27, 2024

[BOUNTY - $500] Llama.cpp inference engine #167

Open

[BOUNTY - $200] Windows native support #186

Open

risingsunomi added 12 commits September 3, 2024 19:26

trying to load in weights but transformers/pytorch doesnt allow that …

f983e93

…unless wanting to rebuild the whole model

adding more testing, refining logit selection

d142be0

working split model test, updating class

be8d7fb

working on class and inference engine updates

9d1ecdd

building out inference engine test

4b0df06

adding working tests, update to forward function to just use input_id…

623468c

…s, update infer to pass logits and hidden_states

cleaning up code and tests, debugging and adding in cleaned up loggin…

19b322d

…g, added model value options

getting infer and stop token issues

cc2c14c

add tracking of next token and other logits into the full input_ids s…

583629c

…o it will stop, context was also dropping due to not having all the logits

grpc testing

7ec5bb8

grpc testing

5903e63

grpc testing

e7a3fd0

risingsunomi added 27 commits November 24, 2024 05:47

adding detach

84f4131

debug process issue

596c715

putting back hidden state pass

a5eb1be

Merge branch 'pr139-dev-oct24' of github.com:risingsunomi/exo-nvidia …

e550af2

…into pr139-dev-oct24

fixing torch inference engine selection not working when adding more …

21e626e

…nodes

fixing typo

e8f689c

fixing llama3 tests, removing mask and input_ids for going through mo…

d0cc3b0

…del layers as that will not be passed to other nodes

optional caching as cache might not work with how sharding works

5f085dc

fix cache assignment

f667735

fixing, set cache to false for inference for now

cb847e4

model double loading vram issue

8325975

model double loading vram issue

c8308b8

setting class shard

2dfce95

working on inference engine issue of too much vram

39cfbf5

fixing vram issue from total_response_length being set with max_seq_l…

185502a

…en and not max_new_tokens

fixing split and full model generation, finetune for nodes and genera…

738e931

…tion, working on exo inference engine

updated full test to generate to stop or max tokens for testing, upda…

907ba0b

…ting split model test, trying to get inference engine to hit stop token

Merge branch 'main' of github.com:exo-explore/exo into exo-explore-main

f3c868b

Merge branch 'exo-explore-main'

f0c8fb1

Merge branch 'main' of github.com:risingsunomi/exo-nvidia

c6806f9

Merge branch 'main' of github.com:risingsunomi/exo-nvidia into pr139-…

4c93855

…dev-oct24

Merge branch 'main' of github.com:risingsunomi/exo-nvidia into pr139-…

b538cd2

…dev-oct24

fixing inference sampling

8c29d27

changing back temp and top_k passing

80122a0

moving back to unsloth llama version for 3.2-1B

3e0d117

Merge pull request #34 from risingsunomi/pr139-dev-oct24

34ca2ad

Pr139 dev oct24

cleaning up code, doing more testing as some bugs a bit still

28d9900

having a check for llama or Llama for loading tensors, will add suppo…

30651ea

…rt for more models like Qwen as building Qwen engine, updated full llama test with more complex prompts

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bounty] PyTorch & HuggingFace Interface #139

[Bounty] PyTorch & HuggingFace Interface #139

risingsunomi commented Aug 11, 2024

AlexCheema commented Aug 18, 2024

risingsunomi commented Aug 18, 2024

AlexCheema Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Oct 2, 2024

risingsunomi Oct 2, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi Aug 23, 2024

dtnewman Oct 11, 2024

AlexCheema Oct 11, 2024

AlexCheema commented Aug 23, 2024

AlexCheema commented Aug 23, 2024

AlexCheema Aug 23, 2024

AlexCheema Aug 23, 2024

AlexCheema Aug 23, 2024

risingsunomi left a comment

risingsunomi commented Dec 1, 2024

risingsunomi commented Dec 1, 2024


		layers.append(layer)

		self.full_model.model.layers = nn.ModuleList(layers)



		# Load the model
		self.full_model = AutoModelForCausalLM.from_pretrained(

[Bounty] PyTorch & HuggingFace Interface #139

Are you sure you want to change the base?

[Bounty] PyTorch & HuggingFace Interface #139

Conversation

risingsunomi commented Aug 11, 2024

What's New?

Limitations and Bugs

Request for Feedback

Important Note on Meta LLaMA 3.1 Model

Chat API Update

AlexCheema commented Aug 18, 2024

risingsunomi commented Aug 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AlexCheema commented Aug 23, 2024

AlexCheema commented Aug 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

risingsunomi left a comment

Choose a reason for hiding this comment

risingsunomi commented Dec 1, 2024

risingsunomi commented Dec 1, 2024