enhance auto device map and support XPU

[xin3he]

Related issue:

• sdpa backward kernel is required to reduce memory usage torch-xpu-ops#2232
• The automatic device map should leverage all GPU memory provided #917
• 70B W4A16 OOM on BMG #712
• support xpu set device or logger warning #557

Env:

4 cards 24G Intel GPU

Llama 8B MXFP4 (w/o torch.compie)

auto-round --model meta-llama/Llama-3.1-8B-Instruct --scheme MXFP4 --device 0,1,2,3 --low_gpu_mem_usage

Auto device map result:

{'self_attn.q_proj': 'xpu:2', 'self_attn.k_proj': 'xpu:2', 'self_attn.v_proj': 'xpu:1', 'self_attn.o_proj': 'xpu:1', 'mlp.gate_proj': 'xpu:1', 'mlp.up_proj': 'xpu:2', 'mlp.down_proj': 'xpu:3'}

card_idx: max_memory_used

0: 19G, 1: 16G; 2: 13G; 3: 21G

Llama 8B W4A16 (w/o torch.compie)

auto-round --model meta-llama/Llama-3.1-8B-Instruct --scheme W4A16 --device 0,1 --low_gpu_mem_usage

Auto device map result:

{'self_attn.q_proj': 'xpu:1', 'self_attn.k_proj': 'xpu:0', 'self_attn.v_proj': 'xpu:0', 'self_attn.o_proj': 'xpu:1', 'mlp.gate_proj': 'xpu:1', 'mlp.up_proj': 'xpu:1', 'mlp.down_proj': 'xpu:1'}

card_idx: max_memory_used

0: 19G, 1: 8G

Device map logic

estimate_tuning_block_mem function output:

• param_memory considers input activation if enable_act_quant=True
• output_memory = pick_samples * seq_len * out_features * element_size
• all memory consider grad by multiply 2.

layer_memory_dict = {
    'q_proj': {'param_memory': 0.018, 'output_memory': 0.024},
    'k_proj': {'param_memory': 0.018, 'output_memory': 0.024}, 
    'v_proj': {'param_memory': 0.018, 'output_memory': 0.024},
    'o_proj': {'param_memory': 0.018, 'output_memory': 0.008}
}  
input_output_memory = 4  # GB
additional_memory = 1  # GB (or 19 for XPU)

card_0_left_memory

card_0_left_memory = device_0_memory - input_output_memory - additional_memory - total_block_output_memory

_allocate_layers_to_devices function:

Example:
    Input:
        device_memory = {"cuda:0": 30.0, "cuda:1": 40.0, "cuda:2": 40.0}
        layer_memory_dict = {
            "q_proj": {"param_memory": 4.0}, "k_proj": {"param_memory": 1.0},
            "v_proj": {"param_memory": 1.0}, "o_proj": {"param_memory": 4.0},
            "gate_proj": {"param_memory": 11.0}, "up_proj": {"param_memory": 11.0},
            "down_proj": {"param_memory": 11.0}
        }
        mem_per_param = 2.0

    Result (allocation order by size):
        1. gate_proj (22GB) -> cuda:2 (largest, prefer last device)
        2. up_proj (22GB) -> cuda:1 (2nd largest, prefer 2nd last device)
        3. down_proj (22GB) -> cuda:0 (3rd largest, cuda:0 has 30GB available)
        4. q_proj (8GB) -> cuda:2 (neighbor of gate_proj, continuity bonus)
        5. o_proj (8GB) -> cuda:2 (neighbor of q_proj, continuity bonus)
        6. k_proj (2GB) -> cuda:1 (neighbor of q_proj via original order)
        7. v_proj (2GB) -> cuda:1 (neighbor of k_proj, continuity bonus)

[xin3he]

FYI. 70B is still OOM with 4 cards due to intel/torch-xpu-ops#2232

[wenhuach21]

FYI. 70B is still OOM with 4 cards due to intel/torch-xpu-ops#2232

70B on cuda costs less than 50G, we may have a deep dive later, could you manually set the devcie_map yourself to move less parameters on you oom card

[yiliu30]

Llama 8B W4A16 (w/o torch.compie)
auto-round --model meta-llama/Llama-3.1-8B-Instruct --scheme W4A16 --device 0,1 --low_gpu_mem_usage

Is it possible to enable torch.compile?

[xin3he]

Thank you, @wenhuach21 & @yiliu30
For your concerns, I updated as below:

• cuda_xpu_device -> gpu_devices
• v tensor is considered in param memory.
• self.pick_samples = self.batch_size * self.gradient_accumulate_steps
• SDPA (attention activation) memory is considers, 1GB for CUDA and 19GB for XPU
• for 70B, the OOM is due to SDPA issue, 19GB is assumed from 8B and it could be larger for 70B due to big hidden dim. no quantized layer are on xpu:0
• torch.compile could reduce the memory usage. I didn't attach it. It works well.

[wenhuach21]

Thank you, @wenhuach21 & @yiliu30 For your concerns, I updated as below:

• cuda_xpu_device -> gpu_devices
• v tensor is considered in param memory.
• self.pick_samples = self.batch_size * self.gradient_accumulate_steps
• SDPA (attention activation) memory is considers, 1GB for CUDA and 19GB for XPU
• for 70B, the OOM is due to SDPA issue, 19GB is assumed from 8B and it could be larger for 70B due to big hidden dim. no quantized layer are on xpu:0
• torch.compile could reduce the memory usage. I didn't attach it. It works well.

in tuning, typically we use micro batch size and global batch size , so pick_samples->global_batch_size?

[xin3he]

@wenhuach21 Agreed.

• pick_samples->self.batch_size
• fix device_map bug #954
• revert self.pick_samples change

[xin3he]

New changes:

• Observe the real memory usage on CUDA, and roughly estimate additional activation memory from SiLU and Norm by doubling layer outputs memory.
• pass output_device into set_auto_device_map_for_block_with_tuning
• add self.low_gpu_mem_usage as constrain

[xin3he]

Cuda:0 memory used with no linear layers on cuda:0.

llama 8B W4A16

• about 6GB on CUDA

Qwen3 8B

• about 6GB on CUDA

Qwen3 32B

• about 10GB on CUDA

Now the estimation is matching the observation.

[xin3he]

BTW, By setting output_device="cpu", which means that only quantized layers are on XPU.
8B could work with less than 10GB memory usage, but the quantization time increase from 3min to 17min. (roughly estimation).
70B could also work, no estimation time.

[wenhuach21]

please have a double check.
Run the default command (auto-round --model xxx)for the 8B model to check whether the speed and accuracy are typically same with the main branch.