AssertionError: Torch not compiled with CUDA enabled when using device_map="auto" in Ascend NPU

[jiaqiw09]

System Info

Ascend NPU
transformers>=4.50.0
torch 2.1

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

When using device_map with Ascend NPU devices in transformers >=4.50.0, loading models fails with assertion errors. The issue occurs because the new loading implementation in _load_state_dict_into_meta_model doesn't properly handle integer device indices for NPU devices, whereas previous versions (<4.50.0) used accelerate.utils.set_module_tensor_to_device which correctly converts integer indices to device strings like "npu:0".

On an Ascend NPU system, attempt to load a model with device mapping:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "gpt2",
    device_map="auto",  # Or custom device_map with integer indices
    torch_dtype=torch.float16
)

Observe the failure with stack trace pointing to modeling_utils.py in _load_state_dict_into_meta_model

Expected behavior

In transformers <=4.49.0, device mapping used accelerate.utils.set_module_tensor_to_device for various device types

In transformers >=4.50.0, the new _load_state_dict_into_meta_model directly uses device values from device_map without converting integer indices to device-specific strings

For NPU devices, integer indices (like 0) are not automatically converted to proper device strings ("npu:0")

[Rocketknight1]

cc @SunMarc

[SunMarc]

In transformers >=4.50.0, the new _load_state_dict_into_meta_model directly uses device values from device_map without converting integer indices to device-specific strings

This might indeed be an issue. Would you like to fix this in transformers by adding this mapping @jiaqiw09 ?

[jiaqiw09]

In transformers >=4.50.0, the new _load_state_dict_into_meta_model directly uses device values from device_map without converting integer indices to device-specific strings

This might indeed be an issue. Would you like to fix this in transformers by adding this mapping @jiaqiw09 ?

@SunMarc Thanks for the suggestion. I checked the code, and I believe the most appropriate place to fix this is here:

https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L4365

        # change device_map into a map if we passed an int, a str or a torch.device
        if isinstance(device_map, torch.device):
            device_map = {"": device_map}
        elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]:
            try:
                device_map = {"": torch.device(device_map)}
            except RuntimeError:
                raise ValueError(
                    "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or "
                    f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}."
                )
        elif isinstance(device_map, int):
            if device_map < 0:
                raise ValueError(
                    "You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' "
                )
            else:
                device_map = {"": device_map}

My plan is to modify the elif isinstance(device_map, int): block, and convert the integer index into a proper device string before wrapping it into a dictionary. For example:

def get_current_device(device_index: Union[str, int]):
    r"""
    Return current device.
    """
    state = PartialState()
    device_type = state.device.type.lower()
    device = f"{device_type}:{device_index}"
    return device

I haven’t found a ready-to-use function in Accelerate that provides this functionality, so I think we may need to handle this manually. However, I'm aware that this could potentially touch a broader range of logic, so I'd appreciate any feedback or suggestions on the best way to integrate this safely.

[SunMarc]

I think the best place would be to add it in _load_state_dict_into_meta_model.

def get_current_device(device_index: Union[str, int]):
    r"""
    Return current device.
    """
    state = PartialState()
    device_type = state.device.type.lower()
    device = f"{device_type}:{device_index}"
    return device

Let's not use PartialState, just copy paste the relevant part in accelerate lib and put it in transformers !

[gspeter-max]

PyTorch Build (NPU-ONLY):
The user has a PyTorch build specifically compiled for Ascend NPUs.
Crucially, this build does not include CUDA libraries. It literally cannot use NVIDIA GPUs.
accelerate's Detection:
When device_map="auto" runs (powered by accelerate), it does correctly detect the presence of NPUs (torch.npu.is_available() is True).
It then decides to map certain layers to NPU 0, NPU 1, etc.
It expresses this decision by placing the integer 0 or 1 as a value in the device_map dictionary.
The "Misinterpretation" at the torch.Tensor.to() Level (The Bug's Cause):
Now, _load_state_dict_into_meta_model receives this device_map dictionary, extracts the integer 0 for a specific parameter (param_device = 0).
Later, it calls param.to(param_device) (which is param.to(0)).
This is the critical point: When PyTorch's low-level C++ code receives a bare integer 0 (or 1, etc.) as a device argument, its default lookup mechanism often prioritizes CUDA if the PyTorch library itself was designed to support CUDA, even if it's not present or enabled in the current build.
Even if the specific PyTorch binary doesn't have CUDA libraries, the internal code path might still try to go: "Is 0 a CUDA device? Let me check a CUDA-related function/variable."
When it tries to call a CUDA-related function (e.g., cudaGetDeviceCount() or cudaMalloc()) that simply does not exist in this NPU-compiled PyTorch binary, that's when the AssertionError: Torch not compiled with CUDA enabled is triggered.

[gspeter-max]

The Bug's Symptom as Proof:

The fact that you get a CUDA error on an NPU system (when you pass 0) is direct evidence that PyTorch's default bare-integer resolution prioritizes or defaults to a CUDA check first, even when NPU is available. If it had implicitly gone "Oh, NPU is available, so 0 means npu:0," then you would never see the "CUDA enabled" error on an NPU-only system.

[gspeter-max]

'''

my fix:

if isinstance(param_device, int): # param_device is the integer 0 from device_map="auto"
    device_id_int = param_device

    # We (Transformers library) explicitly check for NPU first!
    if hasattr(torch, 'npu') and torch.npu.is_available() and device_id_int < torch.npu.device_count():
        param_device = f"npu:{device_id_int}" # We CONVERT 0 to "npu:0"
    # ... (then we check for CUDA, CPU, etc.)

 Later, in the original Transformers code:
 _load_parameter_into_model(model, param_name, param.to(param_device))

By explicitly converting 0 to "npu:0" before calling param.to(), we feed PyTorch the unambiguous string it needs to correctly target the NPU, bypassing its problematic default integer-resolution logic.'''