Add LLAMA

[msaroufim]

Fixes #1443

Notable things I had to do - @ezyang hopefully these are OK

• Since LLAMA requires special permission to download weights and checkpoints and the tokenizer, I went ahead with random checkpoints and random tokenizer - not sure CI qualifies as a valid research endeavour
• I removed the dependency on fairscale so had to make a few adjustments like turning ParallelLinear into Linear or ParallelEmbedding into Embedding and things mostly seem to work fine. And added bonus is you can run the example on a single machine
• Inference in the code using torch inference mode, I removed it since it has a weird interaction with torch.compile
• The open source LLAMA repo is inference only so there is no training support in this script

Some other things I can improve in an another PR

• Better configuration including sequence length and batching
• Reenabling distributed support with FAIRSCALE

I can run the code now

(bench) ubuntu@ip-172-31-39-186:~/benchmark$ python run.py llama -d cuda
Running eval method from llama on cuda in eager mode with input batch size 32.
GPU Time:             10.006 milliseconds
CPU Total Wall Time:  10.045 milliseconds

[xuzhao9]

I get the following error when running the model locally with run.py:

$ python run.py llama -d cpu
Warning: Could not find dependent module llama for Model llama, skip it.
 No module named 'llama'
Unable to find model matching llama.

Can you help try and see if you can reproduce?

[ezyang]

The limitations look fine; I prefer not downloading the checkpoint.

[msaroufim]

Ok this should work

(bench) ubuntu@ip-172-31-38-220:~/benchmark$ python run.py llama -d cpu
Running eval method from llama on cpu in eager mode with input batch size 32.
CPU Total Wall Time:  10.136 milliseconds
CPU Peak Memory:                2.4590 GB
(bench) ubuntu@ip-172-31-38-220:~/benchmark$ python run.py llama -d cuda
Running eval method from llama on cuda in eager mode with input batch size 32.
GPU Time:              8.040 milliseconds
CPU Total Wall Time:   8.067 milliseconds
GPU 0 Peak Memory:              1.9117 GB
CPU Peak Memory:                2.0645 GB

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[xuzhao9]

LGTM!

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[xuzhao9]

I am wondering why torch.set_default_device('cuda') is needed?

[msaroufim]

It's convenient in general to make sure everything runs on the same device that way you dont need to add to(device) calls for inputs and models - but it works for any device so just made that clearer now

[xuzhao9]

Thanks!

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[xuzhao9]

I am curious why we don't need to explicitly move self.model and self.example_inputs to the device here? For example:

self.model = Transformer(self.model_args).to(self.device)
self.example_inputs = (torch.tensor([[1, 1], [1,1]], dtype=torch.int).to(self.device), 1)

[msaroufim]

Fixed

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@xuzhao9 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@xuzhao9 merged this pull request in c78f1f3.

[msaroufim]

This weeekend the internet was buzzing with news of this cpp llama implementation ggerganov/llama.cpp#33 (comment) - It seems like it might not be that tricky to also validate correctness for this model by loading in real checkpoints and tokenizer despite minor differences in the model

Might need to do a round 2 pass on this PR soon