Repository: huggingface/diffusers
File: src/diffusers/models/model_loading_utils.py (Python)

Code before:
    """
    This function warm-ups the caching allocator based on the size of the model tensors that will reside on each
    device. It allows to have one large call to Malloc, instead of recursively calling it later when loading the model,
    which is actually the loading speed bottleneck. Calling this function allows to cut the model loading time by a
    very large margin.
    """
    factor = 2 if hf_quantizer is None else hf_quantizer.get_cuda_warm_up_factor()

    # Keep only accelerator devices
    accelerator_device_map = {
        param: torch.device(device)
        for param, device in expanded_device_map.items()
        if str(device) not in ["cpu", "disk"]
    }
    if not accelerator_device_map:
        return

    elements_per_device = defaultdict(int)
    for param_name, device in accelerator_device_map.items():
        try:
            p = model.get_parameter(param_name)
        except AttributeError:
            try:
                p = model.get_buffer(param_name)
            except AttributeError:
                raise AttributeError(f"Parameter {param_name} not found in model")
        # TODO: account for TP when needed.
        elements_per_device[device] += p.numel()

    # This will kick off the caching allocator to avoid having to Malloc afterwards
    for device, elem_count in elements_per_device.items():
        warmup_elems = max(1, elem_count // factor)
        _ = torch.empty(warmup_elems, dtype=dtype, device=device, requires_grad=False)


Review comment (suggestion):
```suggestion
                raise AttributeError(f"Parameter or buffer with name={param_name} not found in model")
```

or something similar maybe since we check both params and buffers


Code after:
    """
    This function warm-ups the caching allocator based on the size of the model tensors that will reside on each
    device. It allows to have one large call to Malloc, instead of recursively calling it later when loading the model,
    which is actually the loading speed bottleneck. Calling this function allows to cut the model loading time by a
    very large margin.
    """
    factor = 2 if hf_quantizer is None else hf_quantizer.get_cuda_warm_up_factor()

    # Keep only accelerator devices
    accelerator_device_map = {
        param: torch.device(device)
        for param, device in expanded_device_map.items()
        if str(device) not in ["cpu", "disk"]
    }
    if not accelerator_device_map:
        return

    elements_per_device = defaultdict(int)
    for param_name, device in accelerator_device_map.items():
        try:
            p = model.get_parameter(param_name)
        except AttributeError:
            try:
                p = model.get_buffer(param_name)
            except AttributeError:
                raise AttributeError(f"Parameter or buffer with name={param_name} not found in model")
        # TODO: account for TP when needed.
        elements_per_device[device] += p.numel()

    # This will kick off the caching allocator to avoid having to Malloc afterwards
    for device, elem_count in elements_per_device.items():
        warmup_elems = max(1, elem_count // factor)
        _ = torch.empty(warmup_elems, dtype=dtype, device=device, requires_grad=False)