Infer hardware from model name

[yunfeng-scale]

Pull Request Summary

Infer hardware specs from model name so these fields are optional now

formular used:

dtype_size = 2 (for 16bit float)

min_kv_cache_size =
        2
        * dtype_size
        * config["num_hidden_layers"]
        * config["hidden_size"]
        * config["max_position_embeddings"]
        // (config["num_attention_heads"] // config["num_key_value_heads"])

model_weights_size = dtype_size * model_param_count_b * 1_000_000_000

min_memory_gb = math.ceil((min_kv_cache_size + model_weights_size) / 1_000_000_000 / 0.9)

we hard code the param count for MoE models right now.
we omit some other small weights like embedding layer and MLP layer post transformer
this estimates is mostly correct with some issues on long context window models (investigation TBD)

Test Plan and Usage Guide

added unit tests and created llama 3 8b and codellama endpoint

[yixu34]

Do we have a use case-level test that verifies this behavior? If there was a prior test that exercised infer_hardware_from_model_name, we can probably just route to that one.

[seanshi-scale]

what's the reason we're getting rid of the max_model_len/max_num_batched_tokens args to vllm?

[yunfeng-scale]

i forgot why we're doing this in the first place, but i'm pretty certain in recent version of vLLM this is not needed. also checked most of these models, config.json has the same max_position_embeddings with the values here

[yixu34]

👀

[yixu34]

1. Cleanup prints
2. Unrelated changes?

[seanshi-scale]

probably not really an issue, but how well do these map to Azure instance types?

[yunfeng-scale]

with my limited experiences on Azure, they do provide the same set of GPUs @squeakymouse wdyt?

[seanshi-scale]

guess we're gonna handle quantization later?

[yunfeng-scale]

quantization can already be handled here. but I chose not to update dtype_size since in my experiences, they usually make things slower not faster (at least for bitsandbytes and AWQ), so in order to achieve the same speed we still need the same amount of GPUs

[seanshi-scale]

IIUC we're implicitly setting this to "batch size = 1 and filling up the context window", this feels reasonable, but would it make sense to add a bit of room for a larger batch size? (I guess for something with a shorter context window e.g. llama 2 it makes more sense to add some room, maybe less so for mixtral)

[yunfeng-scale]

for all the existing models i tested, batch size = 1 is a good enough default to reach to reasonable amount of GPUs. model builders would have thought about the tradeoffs between model sizes and GPU sizes

[seanshi-scale]

should we just deprecate the num_shards field at this point?

[yunfeng-scale]

yes we should but would prefer not in this PR

[seanshi-scale]

nit: "... if any hardware spec field is provided"?

[yunfeng-scale]

#515 (comment)
i guess it works both ways. since here i'm only allowing two states: either all fields are provided, or none of them are provided

[seanshi-scale]

had a few questions/nits, but lgtm!