自定义HuggingFace生成函数时Top-p采样触发CUDA断言错误

问题背景 #

为实现HuggingFace Transformers原生model.generate()不支持的logit-bias功能，自行重写模型生成函数，但添加Top-p采样逻辑后触发CUDA设备端断言错误，无Top-p采样时功能正常。

相关代码片段 #

generation_args = {
    "max_new_tokens": 500,
    "temperature": 0.4,
    "do_sample": True,
    "top_p": 0.5,
    "top_k": None,
}


def top_p_filtering(logits, top_p):
    """Filter the logits using top-p (nucleus) sampling."""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_keep = cumulative_probs <= top_p
    sorted_indices_to_keep[..., 0] = True
    logits[sorted_indices[~sorted_indices_to_keep]] = float('-inf')
    return logits


def custom_generate(input_ids, streamer, max_new_tokens, temperature, top_p):
    past_key_values = None
    attention_mask = torch.ones(input_ids.shape, device=input_ids.device)

    for _ in range(max_new_tokens):
        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                past_key_values=past_key_values,
                attention_mask=attention_mask,
                use_cache=True
            )

        logits = outputs.logits[:, -1, :]

        if temperature != 1.0:
            logits = logits / temperature

        if top_p is not None and top_p < 1.0:
            logits = top_p_filtering(logits, top_p)

        next_token_probs = torch.nn.functional.softmax(logits, dim=-1)
        next_token_id = torch.multinomial(next_token_probs, num_samples=1)  
        streamer.put(next_token_id)

        input_ids = next_token_id
        attention_mask = torch.cat(
            [attention_mask, torch.ones((attention_mask.shape[0], 1), device=attention_mask.device)], dim=1)
        past_key_values = outputs.past_key_values

        if next_token_id.item() == tokenizer.eos_token_id:  
            break

with torch.no_grad():
    custom_generate(input_ids, streamer, generation_args["max_new_tokens"], generation_args["temperature"], generation_args["top_p"])

报错信息 #

../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [10,0,0], thread: [63,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
Exception in thread Thread-18 (generate):
Traceback (most recent call last):
  File "/usr/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/c/Users/User/Documents/EmpatheticChatBot/Inference-Server.py", line 130, in generate
    custom_generate(input_ids, streamer, generation_args["max_new_tokens"], generation_args["temperature"], generation_args["top_p"])
  File "/mnt/c/Users/User/Documents/EmpatheticChatBot/Inference-Server.py", line 108, in custom_generate
    next_token_id = torch.multinomial(next_token_probs,
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

错误原因 #

top_p_filtering函数中的索引逻辑存在维度不匹配问题：

• sorted_indices的形状为[batch_size, vocab_size]，~sorted_indices_to_keep是同形状的布尔掩码
• 直接使用sorted_indices[~sorted_indices_to_keep]会得到一个一维数组，包含所有batch中需要过滤的token索引，但这些索引是相对于每个batch内部的vocab维度，而非全局索引
• 用一维索引去赋值二维的logits张量时，会导致索引越界，触发CUDA断言错误

修复方案 #

修改top_p_filtering函数，正确处理每个batch样本的过滤逻辑，确保索引操作维度匹配。以下是两种可行实现：

方案1：逐样本处理（直观易读） #

def top_p_filtering(logits, top_p):
    """Filter the logits using top-p (nucleus) sampling."""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
    cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
    
    # 标记需要保留的token，确保至少保留第一个token
    sorted_indices_to_keep = cumulative_probs <= top_p
    sorted_indices_to_keep[..., 0] = True
    
    # 创建与logits同形状的掩码，逐个batch填充保留的token
    mask = torch.zeros_like(logits, dtype=torch.bool)
    batch_size = logits.shape[0]
    for i in range(batch_size):
        # 对第i个样本，将需要保留的sorted索引对应位置设为True
        mask[i, sorted_indices[i][sorted_indices_to_keep[i]]] = True
    
    # 将不保留的token logits设为负无穷
    logits = torch.where(mask, logits, torch.tensor(float('-inf'), device=logits.device))
    return logits

方案2：向量式处理（高效无循环） #

def top_p_filtering(logits, top_p):
    """Filter the logits using top-p (nucleus) sampling."""
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
    cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
    
    # 找到每个样本中第一个超过top_p的位置，确保至少保留1个token
    cutoff = torch.argmax((cumulative_probs > top_p).float(), dim=-1)
    cutoff = torch.max(cutoff, torch.tensor(1, device=logits.device))
    
    # 创建掩码：每个样本保留前cutoff个排序后的token
    mask = torch.arange(sorted_logits.shape[-1], device=logits.device)[None, :] < cutoff[:, None]
    
    # 对排序后的logits应用掩码，再转回原顺序
    sorted_logits = torch.where(mask, sorted_logits, torch.tensor(float('-inf'), device=logits.device))
    logits = torch.gather(sorted_logits, dim=-1, index=torch.argsort(sorted_indices, dim=-1))
    
    return logits

验证步骤 #

1. 替换原top_p_filtering函数为修复后的版本
2. 运行代码前可添加校验：在next_token_probs = torch.nn.functional.softmax(logits, dim=-1)后，检查是否存在NaN或全0的概率分布：

   assert not torch.isnan(next_token_probs).any(), "存在NaN概率"
   assert (next_token_probs.sum(dim=-1) > 0).all(), "存在全0概率分布"
   

3. 重新运行生成逻辑，确认CUDA断言错误消失

内容的提问来源于stack exchange，提问作者AndreiS