Add DTensor LLaMA inference model: simple_gpt

[xmfan]

Adds simple_gpt + DTensor implemented in https://github.com/pytorch-labs/simple_gpt/pull/7 to torchbench

Tested via python benchmarks/dynamo/torchbench.py -d cuda --output-directory=benchmark_logs --output=performance.csv --inference --performance --timing --print-memory --multiprocess --nothing --only simple_gpt. Note: --nothing is used here to disable compile, since DTensor + compile isn't yet supported in main

dev,name,batch_size,speedup,abs_latency,compilation_latency,compression_ratio,eager_peak_mem,dynamo_peak_mem,calls_captured,unique_graphs,graph_breaks,unique_graph_breaks
cuda,simple_gpt,1,0.966153,196.819773,-0.059319,1.000000,4.576880,4.576880,0,0,0,0
cuda,simple_gpt,1,0.967389,196.608152,-0.058833,1.000000,4.577404,4.577404,0,0,0,0
cuda,simple_gpt,1,0.973152,196.093583,-0.059316,1.000000,4.593133,4.593133,0,0,0,0
cuda,simple_gpt,1,0.973087,196.124046,-0.075580,1.000000,4.611483,4.611483,0,0,0,0
cuda,simple_gpt,1,0.967908,193.998484,-0.040192,1.000000,4.593133,4.593133,0,0,0,0
cuda,simple_gpt,1,0.968949,193.798088,-0.028878,1.000000,4.593133,4.593133,0,0,0,0

2 changes were required to the model:

• decorate torch.no_grad() on the caches, previously this was done outside the model, the entire eval call was wrapped in a torch.no_grad() context. After using torchbench, I notice even with only inference mode, we don't disable gradient calculations
• rank/world size, added support from torchbench side in Forward rank and world size info to Torchbench models when using dynamo runner pytorch#108438 and updated model to fetch from the provided extra_args

[H-Huang]

Looks good

[H-Huang]

Just curious, is the only use of Lightning this init_module? I'm not sure what this does, but can we remove it by initializing with meta device? Then the implementation doesn't rely on third-party libraries can stay as close to native pytorch as possible.

[xmfan]

Yeah good point, let me update this with the latest init code that removed the lightning dep

[Chillee]

I've actually updated the initial code to avoid the lightning dependency (and it's actually much faster too!)

https://github.com/pytorch-labs/simple_gpt/blob/main/generate.py#L162

[xmfan]

For nightly runs, do we usually load the actual weights? The load time should be quick, but the weights file is 10+GB just for LLaMA-7B. Otherwise, I could also update the default weights initialization to use random instead of just torch.zeros

[msaroufim]

Use the real weights otherwise you might fail accuracy checks in unpredictable ways

[xmfan]

Is this still needed if we only run this model through dynamorunner (model requires distributed, multiple GPUs)? The unit tests from this repo are skipped i.e. python test.py -k 'test_simple_gpt_', which includes the accuracy checks

[xuzhao9]

decorate torch.no_grad() on the caches, previously this was done outside the model, the entire eval call was wrapped in a torch.no_grad() context. After using torchbench, I notice even with only inference mode, we don't disable gradient calculations

This is because some models in torchbench (e.g. maml) does need gradient calculations in inference mode. So we leave the choice of whether enabling the gradient context to the eval test code.

There is an open issue about testing that the no_grad is enabled in the eval test: #1838

[xmfan]

@pytorchbot merge

[facebook-github-bot]

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

[facebook-github-bot]

@xmfan merged this pull request in 0a64c55.