Segmentation Fault with Qwen2 on Multi-GPU (device_map='auto')
Author: justice-Pan
Pull request: no
Labels: name: bug

### System Info


**Library versions**
- `transformers`: 4.40.1
- `accelerate`: 0.29.3
- `torch`: 2.1.2+cu121
- `bitsandbytes`: 0.41.3
- `CUDA`: 12.1
- `Python`: 3.10

**Hardware**
- GPUs: 2x NVIDIA RTX 3090 (24GB)
- Driver: NVIDIA 535.171.04


### Who can help?


**Description**
When attempting to load the `deepseek-ai/DeepSeek-R1-Distill-Qwen-14B` model (a `qwen2` model type) using `device_map="auto"` on a multi-GPU system, a `Segmentation fault (core dumped)` occurs immediately after the model weights are loaded.

**The bug persists even when 4-bit quantization (`bitsandbytes`) and PEFT (LoRA) are completely disabled.** This strongly suggests the issue is not with those libraries, but with a core incompatibility between the Qwen2 model's remote code implementation and `accelerate`'s `device_map` functionality.

**Minimal Reproducible Code**
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# This is sufficient to trigger the crash
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16 # Using float16 to even attempt loading
)
print("Model loading finished. The script crashes right after this.")
# The script never reaches this point
```

**Expected behavior**
The script should either load the model successfully or fail with a clear `CUDA out of memory` error if the model is too large.

**Actual behavior**
The script prints `Loading checkpoint shards... 100%` and then immediately exits with `Segmentation fault (core dumped)`.

This bug prevents any form of multi-GPU inference or training with Qwen2 models using `device_map` on this hardware configuration.

### 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

import os
import yaml
import argparse
import torch
import warnings
import importlib.metadata

# 忽略一些不影响结果的警告
warnings.filterwarnings("ignore")

from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)

def main(config_path: str):
    # --- 1. 加载配置 ---
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    model_cfg = config['model_params']
    data_cfg = config['data_params']
    lora_cfg = config['lora_params']
    training_cfg = config['training_params']
    logging_cfg = config.get('logging_params', {})
    optim_cfg = config.get('optim_params', {})

    print("--- 配置加载完毕 ---")

    # --- 2. 检查并升级关键库版本 ---
    print("--- 正在检查关键库版本... ---")
    try:
        # 使用 importlib.metadata 进行更可靠的版本检查
        bitsandbytes_version = importlib.metadata.version("bitsandbytes")
        peft_version = importlib.metadata.version("peft")
        accelerate_version = importlib.metadata.version("accelerate")
        
        print(f"bitsandbytes version: {bitsandbytes_version}")
        print(f"peft version: {peft_version}")
        print(f"accelerate version: {accelerate_version}")
        
    except importlib.metadata.PackageNotFoundError as e:
        print(f"错误: 无法找到包的版本: {e}。请确保库已在当前环境中正确安装。")
        return
    except Exception as e:
        print(f"检查版本时发生未知错误: {e}")
        return
    # --- 3. 加载Tokenizer ---
    tokenizer = AutoTokenizer.from_pretrained(
        model_cfg['model_dir'],
        trust_remote_code=True,
        local_files_only=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    print("--- Tokenizer加载完毕 ---")

    # --- 4. [关键修改] 加载4-bit量化模型 ---
    # 使用fp16作为计算类型，这在3090上更稳定
    compute_dtype = torch.float16
    
    # bnb_config = BitsAndBytesConfig(
    #     load_in_4bit=True,
    #     bnb_4bit_quant_type="nf4",
    #     bnb_4bit_compute_dtype=compute_dtype,
    #     bnb_4bit_use_double_quant=True,
    # )

    print("--- 正在加载模型 (已禁用量化，预期会OOM)... ---")
    model = AutoModelForCausalLM.from_pretrained(
        model_cfg['model_dir'],
        # quantization_config=bnb_config,
        device_map="auto", # 保持模型并行
        trust_remote_code=True,
        local_files_only=True,
        attn_implementation="sdpa"
        # attn_implementation="flash_attention_2"
        # use_flash_attention_2=True, # [关键新增] 使用Flash Attention 2节省显存并加速
    )
    
    # # 修复一个已知的bug，当使用梯度检查点时需要设置
    # model.config.pretraining_tp = 1

    # model = prepare_model_for_kbit_training(model)
    
    # peft_config = LoraConfig(**lora_cfg, task_type="CAUSAL_LM")
    # model = get_peft_model(model, peft_config)
    # model.print_trainable_parameters()
    # print("--- 4-bit量化与LoRA模型准备完毕 ---")

    # --- 5. 加载和预处理数据集 ---
    def preprocess_function(examples):
        # 假设数据格式是 {"messages": [...]}
        prompts = [tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False) for msgs in examples["messages"]]
        inputs = tokenizer(prompts, max_length=data_cfg['max_seq_length'], truncation=True, padding=False)
        # 标签也需要被截断
        labels = tokenizer(prompts, max_length=data_cfg['max_seq_length'], truncation=True, padding=False)
        inputs['labels'] = labels['input_ids']
        return inputs

    dataset = load_dataset("json", data_files=data_cfg['train_file'], split="train")
    processed_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
    print("--- 数据集处理完毕 ---")

    # --- 6. [关键修改] 配置训练参数 ---
    trainer_args = TrainingArguments(
        output_dir=model_cfg['output_dir'],
        per_device_train_batch_size=training_cfg['per_device_train_batch_size'],
        gradient_accumulation_steps=training_cfg['gradient_accumulation_steps'],
        
        # 使用8位分页优化器
        optim=optim_cfg.get('optim', 'paged_adamw_8bit'),
        
        learning_rate=training_cfg['learning_rate'],
        weight_decay=training_cfg['weight_decay'],
        num_train_epochs=training_cfg['num_train_epochs'],
        lr_scheduler_type=training_cfg['lr_scheduler_type'],
        warmup_ratio=training_cfg['warmup_ratio'],
        max_grad_norm=training_cfg['max_grad_norm'],
        
        # 混合精度与梯度检查点
        fp16=True, # 明确使用fp16
        bf16=False,
        gradient_checkpointing=training_cfg['use_gradient_checkpointing'],
        
        # 日志与保存
        logging_steps=logging_cfg.get('logging_steps', 10),
        save_strategy=logging_cfg.get('save_strategy', 'epoch'),
        save_steps=logging_cfg.get('save_steps', 500),
        save_total_limit=2,
        
        # 禁用tf32，有时它会与4-bit量化冲突
        tf32=False, 
        
        report_to="wandb" if logging_cfg.get('wandb_project') else "none",
        run_name=f"qwen-qlora-{os.path.basename(model_cfg['output_dir'])}"
    )

    # --- 7. 初始化并开始训练 ---
    trainer = Trainer(
        model=model,
        args=trainer_args,
        train_dataset=processed_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True),
    )
    
    # 修复 `use_cache=True` 与梯度检查点的冲突
    if training_cfg['use_gradient_checkpointing']:
        model.config.use_cache = False

    print("--- 开始训练 ---")
    trainer.train()

    # --- 8. 保存最终模型 ---
    print("--- 训练完成，正在保存最终模型... ---")
    trainer.save_model(os.path.join(model_cfg['output_dir'], "final_checkpoint"))
    print("--- 最终模型已保存 ---")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str, required=True, help="Path to the config.yaml file")
    args = parser.parse_args()
    main(args.config)

### Expected behavior

---
name: "Bug Report: Segmentation Fault with Qwen2 on Multi-GPU (device_map='auto')"
---

**Library versions**
- `transformers`: 4.40.1
- `accelerate`: 0.29.3
- `torch`: 2.1.2+cu121
- `bitsandbytes`: 0.41.3
- `CUDA`: 12.1
- `Python`: 3.10

**Hardware**
- GPUs: 2x NVIDIA RTX 3090 (24GB)
- Driver: NVIDIA 535.171.04

**Description**
When attempting to load the `deepseek-ai/DeepSeek-R1-Distill-Qwen-14B` model (a `qwen2` model type) using `device_map="auto"` on a multi-GPU system, a `Segmentation fault (core dumped)` occurs immediately after the model weights are loaded.

**The bug persists even when 4-bit quantization (`bitsandbytes`) and PEFT (LoRA) are completely disabled.** This strongly suggests the issue is not with those libraries, but with a core incompatibility between the Qwen2 model's remote code implementation and `accelerate`'s `device_map` functionality.

**Minimal Reproducible Code**
```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# This is sufficient to trigger the crash
model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.float16 # Using float16 to even attempt loading
)
print("Model loading finished. The script crashes right after this.")
# The script never reaches this point
```

**Expected behavior**
The script should either load the model successfully or fail with a clear `CUDA out of memory` error if the model is too large.

**Actual behavior**
The script prints `Loading checkpoint shards... 100%` and then immediately exits with `Segmentation fault (core dumped)`.

This bug prevents any form of multi-GPU inference or training with Qwen2 models using `device_map` on this hardware configuration.

Comments:
hi @justice-pan, is `trust_remote_code=True` necessary here? Qwen-14B models should be supported in the main library.

--------------原始邮件--------------
发件人："Matt ***@***.***&gt;;
发送时间：2025年10月2日(星期四) 晚上8:06
收件人："huggingface/transformers" ***@***.***&gt;;
抄送："潘浩楠 ***@***.***&gt;;"Mention ***@***.***&gt;;
主题：Re: [huggingface/transformers] Segmentation Fault with Qwen2 on Multi-GPU (device_map='auto') (Issue #41278)
-----------------------------------

 Rocketknight1 left a comment (huggingface/transformers#41278)
 
hi @justice-Pan, is trust_remote_code=True necessary here? Qwen-14B models should be supported in the main library.
 
—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***&gt;

I can not directly download the model from huggingface using the API due to firewalls (I am in China). I need to first download the model onto the server first. That is why I need to use&nbsp;trust_remote_code=True.

Thank you for your replying!

The reason I use trust_remote_code=True is that I’m in mainland China and cannot access HuggingFace’s API to retrieve LLM model parameters and settings. Instead, I have to download them locally on the server in advance and then read the files to obtain the LLM configuration.

--------------原始邮件--------------
发件人："Matt ***@***.***&gt;;
发送时间：2025年10月2日(星期四) 晚上8:06
收件人："huggingface/transformers" ***@***.***&gt;;
抄送："潘浩楠 ***@***.***&gt;;"Mention ***@***.***&gt;;
主题：Re: [huggingface/transformers] Segmentation Fault with Qwen2 on Multi-GPU (device_map='auto') (Issue #41278)
-----------------------------------

 Rocketknight1 left a comment (huggingface/transformers#41278)
 
hi @justice-Pan, is trust_remote_code=True necessary here? Qwen-14B models should be supported in the main library.
 
—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***&gt;

Hello,

Thank you for the quick response. That's a very important clarification.

I need to point out that I am&nbsp;not&nbsp;using an official&nbsp;Qwen/&nbsp;model. I am using a specific distilled version from another organization:&nbsp;deepseek-ai/DeepSeek-R1-Distill-Qwen-14B.

Given that this is a third-party, modified version, I believe&nbsp;trust_remote_code=True&nbsp;is indeed necessary to load its custom architecture, which likely resides within its own repository.

The&nbsp;Segmentation fault&nbsp;occurs even when I disable all quantization and LoRA features, but keep&nbsp;device_map="auto"&nbsp;on a multi-GPU setup (2x RTX 3090). This leads me to believe the bug is in the interaction between this specific model's remote code and the&nbsp;accelerate&nbsp;library's model parallelism logic.

Could you please check if this clarification helps narrow down the problem?

Best regards, justice-Pan

--------------原始邮件--------------
发件人："Matt ***@***.***&gt;;
发送时间：2025年10月2日(星期四) 晚上8:06
收件人："huggingface/transformers" ***@***.***&gt;;
抄送："潘浩楠 ***@***.***&gt;;"Mention ***@***.***&gt;;
主题：Re: [huggingface/transformers] Segmentation Fault with Qwen2 on Multi-GPU (device_map='auto') (Issue #41278)
-----------------------------------

 Rocketknight1 left a comment (huggingface/transformers#41278)
 
hi @justice-Pan, is trust_remote_code=True necessary here? Qwen-14B models should be supported in the main library.
 
—
Reply to this email directly, view it on GitHub, or unsubscribe.
You are receiving this because you were mentioned.Message ID: ***@***.***&gt;

Hi @justice-pan, you shouldn't need `trust_remote_code` in that case, even if you copy it to your local machine! The rule is that it's only necessary when models include their own code (in a file like `modeling_qwen.py`). Deepseek-R1-Distill-Qwen-14B is a standard Qwen format, and so it can be loaded with library code.

However, that probably doesn't resolve the issue you were getting. I would guess that `torch.float16` is a possible culprit. Very few models run well with `float16` as the model dtype. Instead, I would try  `torch.bfloat16` or quantizing with `bitsandbytes` without setting `dtype=torch.float16`. `float16` is generally a legacy format at this point, but `bfloat16` is very common and well-supported.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.