Extend usability of calculate_offload_device_map

[kylesayrs]

Purpose

• Allow calculate_offload_device_map to be used in environments with non-homogenous and/or non-sequential GPUs

Changes

• Default to using all gpus if num_gpus is not specified
• Add gpu_ids argument to allow users to choose which devices to use
• Fix gpu memory calculation in non-homogenous GPU set-ups
  • Previously the calculation assumed that the first GPU's memory was the same for all gpus

Testing

• Grepped for all uses of function and confirmed that changes are compatible
• TODO: Test with simulated non-homogenous GPU setup

[kylesayrs]

In hindsight, I think I'd prefer to give users helper functions which they can use to compute their own device maps. For example

from llmcompressor import hessian_memory_requirements, quantization_memory_requirements, batch_memory_requirements

model_skeleton = load_model_skeleton(model_stub)
reserved_memory = (
    hessian_memory_requirements(model_skeleton) +
    quantization_memory_requirement(model_skeleton) +
    batch_memory_requirements((bs, seq_len), attention_mask=False) +
    ("whatever junk or padding the user thinks is relevant")
)

device_map = infer_auto_device_map(
    model_skeleton,
    max_memory=get_max_memory(reserved_memory, gpu_ids=[1, 2]),
    no_split_module_classes=model_skeleton._no_split_modules,
)

device_map = get_uniform_device_map(model_skeleton, reserved_memory, gpu_ids=[1, 2])

I believe that this is a preferable user experience as opposed to trying to hide too many things behind a function api which a user has to learn.

[dsikka]

This is somewhat difficult to review atm.

Can you summarize what each of these helper functions you're suggesting to use and how the interface is expected to change before and after?

Generally speaking, having the helper functions is nice but we should maintain a higher level api that most users can just use/does the necessary memory calculations for them, which for most people right now would not include batching memory (but we could always expand to include this)