[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

matching main

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

merge fork

merge with fork main

…nsor to numpy

…ference_engine.py, adding in pytorch option for inference engine

… along with adding in torch topk, added in better random distribution selection for when top_p is too low or high, started work on forward_layer_cached but infer functions need to be changed and take any and not a string

…on from and to json

…g weights properly

…ention but not implemented fully, added RMSNorm from modeling llama on HF, added weight renaming and loading along with handling no lm_head weight in safetensor where you then use embed weight as seen with gpt2, still not generating proper reponses further dev being done

risingsunomi · 2024-11-03T11:23:29Z

Finding building a pure pytorch implementation isn't working even from all the examples. Going to try to use the official meta code and hack it for the sharding we need. Will keep trying my method but not making much progress as able to shard the safetensors and everything but inference is not working at all. Still hitting at it.

Any other eyes to look at this would be appreciated. Right now, its in shambles but I am using the torchtune method as opposed to using fairscale. Think I might switch to fairscale though as the official meta llama model is looking better

My WIP code

Sorry again on the delay for this as regular job has me swamped but going to try to hit this faster before month is out.

Thank you again

risingsunomi · 2024-11-03T12:51:46Z

Finding building a pure pytorch implementation isn't working even from all the examples. Going to try to use the official meta code and hack it for the sharding we need. Will keep trying my method but not making much progress as able to shard the safetensors and everything but inference is not working at all. Still hitting at it.

Any other eyes to look at this would be appreciated. Right now, its in shambles but I am using the torchtune method as opposed to using fairscale. Think I might switch to fairscale though as the official meta llama model is looking better

My WIP code

Sorry again on the delay for this as regular job has me swamped but going to try to hit this faster before month is out.

Thank you again

spoke too soon, think I can get this working

…ting and then inference engine testing but we are almost there. HELL YEAAAAAAAAAAA

…ng out next parts for distributed inference

…rking on split modeling

…last layer

…g passed

…it and full test, separating huggingface and torchtune inference engines

…inference is still very high

…, adding hf engine tests

Fork merge

risingsunomi added 2 commits August 10, 2024 18:29

removing broken llama.py

94de83f

Merge pull request #2 from exo-explore/main

f63f4b0

matching main

Merge branch 'main' into main

ab273d3

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

risingsunomi added 11 commits August 23, 2024 17:52

Merge pull request #4 from exo-explore/main

c365749

merge fork

Merge pull request #5 from exo-explore/main

0b8221f

merge with fork main

removing unittest, update inference return type, fixing converting te…

226a0ac

…nsor to numpy

adding nvidia quadro and t1000 support

e11bebd

updating test, updating model selection for smaller quant llama3 model

778cb6e

added updating model options to update_deps.py

56aae50

updating inference class init to take shard, updating pytorch test_in…

7df4640

…ference_engine.py, adding in pytorch option for inference engine

adding updates for inference_engine.py

aa769ca

reducing layer amount for llama3-2b-base

08e8b41

forward rewrite, adding in caching with dynamic cache, cache conversi…

7bcd35e

…on from and to json

dtnewman mentioned this pull request Oct 30, 2024

Package exo as installable #354

Closed

risingsunomi added 5 commits October 30, 2024 03:20

updating attentions, changed model struct, fixing kv cache

ea868c6

fixing kvcache for multiheadattention, fixing layers names for loadin…

f1822e2

…g weights properly

doing work with position_id and causal mask

38028c0

updating torch readme with current model in development

0fd1797

risingsunomi added 22 commits November 3, 2024 04:55

FINALLY A WORKING PYTORCH ONLY MODEL, working on logit gen, shard tes…

b2b63c3

…ting and then inference engine testing but we are almost there. HELL YEAAAAAAAAAAA

cleaning up custom dot product attention but might be removed, buildi…

f53ebd1

…ng out next parts for distributed inference

first layer run fixes, variable layer length weight loading fixes, wo…

e8db8ee

…rking on split modeling

made it so weight for last output layer is only loaded when shard is …

22bc6a7

…last layer

working on sharding issue where hidden state is not working when bein…

7f2abc3

…g passed

fixing last hidden value handling

bdf3240

update test

227199f

update test

5af6302

update test

d7e5aca

update test, turn on caching

1874d23

test safetensor load

3a0ad62

test hidden alignment

6098ae5

updates to torchtune model, fixing non-generation errors, created spl…

fa1e70f

…it and full test, separating huggingface and torchtune inference engines

split model working, updates to safetensor loading letting shard control

d958bf9

reduced model loading ram by loading only some layers in layer list, …

c8bdb09

…inference is still very high

updating readme

75817eb

building out torch inference engine

73630d1

creating torch inference engine, separated torch and hf torch engines…

ad99332

…, adding hf engine tests

Merge github.com:exo-explore/exo into fork-merge

9ff2cc8

merge

6ab6f1c

Merge pull request #32 from risingsunomi/fork-merge

811befc

Fork merge

Merge branch 'exo-explore:main' into main

9f52f24

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 Nov 3, 2024

risingsunomi commented Nov 3, 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 Nov 3, 2024

risingsunomi commented Nov 3, 2024