|
|
import gc |
|
|
import torch |
|
|
import accelerate |
|
|
|
|
|
|
|
|
def get_module_by_name_suffix(model, module_name: str): |
|
|
for name, module in model.named_modules(): |
|
|
if name.endswith(module_name): |
|
|
return module |
|
|
|
|
|
|
|
|
def simple_dispatch_model(model, device_map): |
|
|
from accelerate.hooks import add_hook_to_module, AlignDevicesHook |
|
|
|
|
|
if "" in device_map: |
|
|
d = device_map[""] |
|
|
model = model.to(torch.device(d)) |
|
|
model.hf_device_map = device_map |
|
|
return model |
|
|
|
|
|
tied_params = accelerate.utils.modeling.find_tied_parameters(model) |
|
|
if set(device_map.values()) == {"cpu"} or set(device_map.values()) == { |
|
|
"cpu", |
|
|
"disk", |
|
|
}: |
|
|
main_device = "cpu" |
|
|
else: |
|
|
main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0] |
|
|
|
|
|
cpu_offload_group = [(n, d) for n, d in device_map.items() if d == "cpu"] |
|
|
prev_hook = None |
|
|
for idx, (n, d) in enumerate(cpu_offload_group): |
|
|
m = get_module_by_name_suffix(model, n) |
|
|
_, prev_hook = accelerate.cpu_offload_with_hook( |
|
|
m, execution_device=main_device, prev_module_hook=prev_hook |
|
|
) |
|
|
|
|
|
if len(cpu_offload_group) > 1: |
|
|
get_module_by_name_suffix( |
|
|
model, cpu_offload_group[0][0] |
|
|
)._hf_hook.prev_module_hook = prev_hook |
|
|
|
|
|
for n, d in device_map.items(): |
|
|
m = get_module_by_name_suffix(model, n) |
|
|
if d != "cpu": |
|
|
d = torch.device(d) |
|
|
hook = AlignDevicesHook(d, io_same_device=True, place_submodules=True) |
|
|
add_hook_to_module(m, hook) |
|
|
accelerate.utils.modeling.retie_parameters(model, tied_params) |
|
|
model.hf_device_map = device_map |
|
|
|
|
|
return model |
|
|
|
|
|
|
|
|
def set_module_name(model, name, value): |
|
|
if "." in name: |
|
|
parent_name = name.rsplit(".", 1)[0] |
|
|
child_name = name[len(parent_name) + 1 :] |
|
|
parent = model.get_submodule(parent_name) |
|
|
else: |
|
|
parent_name = "" |
|
|
parent = model |
|
|
child_name = name |
|
|
|
|
|
setattr(parent, child_name, value) |
|
|
|
|
|
|
|
|
def clear_memory(weight=None): |
|
|
if weight is not None: |
|
|
del weight |
|
|
gc.collect() |
|
|
torch.cuda.empty_cache() |
|
|
|
|
|
|
|
|
def compute_memory_used_pct(device): |
|
|
memory_used = torch.cuda.max_memory_allocated(device) / (1024**3) |
|
|
memory_pct = ( |
|
|
memory_used |
|
|
/ (torch.cuda.get_device_properties(device).total_memory / (1024**3)) |
|
|
* 100 |
|
|
) |
|
|
return memory_pct |
|
|
|
|
|
|
|
|
def get_best_device(): |
|
|
if torch.backends.mps.is_available(): |
|
|
return "mps" |
|
|
elif torch.cuda.is_available(): |
|
|
return "cuda:0" |
|
|
else: |
|
|
return "cpu" |
|
|
|
|
|
|
|
|
def get_lowest_memory_device_index(): |
|
|
device = None |
|
|
curr_device_memory_pct = 0 |
|
|
for device_index in range(torch.cuda.device_count()): |
|
|
device_memory_pct = compute_memory_used_pct(device_index) |
|
|
if device is None or device_memory_pct < curr_device_memory_pct: |
|
|
device = device_index |
|
|
curr_device_memory_pct = device_memory_pct |
|
|
|
|
|
return device |
|
|
|
