Lorax server hangs on warming up the server when total_tokens is larger than 4096 with LLama2-7b-hf on a l4 GPU

## The problem
Issue predibase/lorax#334 (closed, 11 comments): ### System Info These are the params I am using to start the lorax server (in GKE): ``` - '--model-id=meta-llama/Llama-2-7b-hf' - '--revision=8cca527612d856d7d32bd94f8103728d614eb852' - '--source=hub' # greater than the max input tokens we will use in the benchmarks - '--max-input-length=1048' # greater than the sum of max output tokens + max input tokens # for any benchmark. The max we have is 512 input tokens + 3584 output tokens - '--max-total-tokens=4096' - '--max-batch-prefill-tokens=2048' # run on a single GPU for now - '--num-shard=1' - '--cuda-memory-fraction=0.

## The fix
[tgaddair (contributor)]: The server hanging during warmup was a real hang, not just slow warmup. Fixed in PR #352; the reporter confirmed the server starts up without issues after the fix. If it only seems stuck, give warmup up to 3 minutes first, since warmup generates max_total_tokens minus max_input_length tokens.

## Notes
Thread: https://github.com/predibase/lorax/issues/334. Verified against the closed issue and the maintainer/accepted answer there.