pytorch LLM训练过程中的精度调试实践

本文记录了,在某加速卡上进行LLM训练,精度问题的定位过程

1.查看权值的最大,最小值

tee dump_weight.py <<-'EOF'
import torch
import json
import numpy as np

weight_files=set()
with open('../llama-2-7b-hf/pytorch_model.bin.index.json', 'r') as f:
    index_data = json.load(f)["weight_map"]
for k,v in index_data.items():
    weight_files.add(v)
gmax=[]
gmin=[]
for i in weight_files:
    for k,w in torch.load(f"../llama-2-7b-hf/{i}",map_location="cpu").items():
        print(f"{k:<64s},max:{w.max().item():8.2f},{w.min().item():8.2f}")
        gmax.append(w.max().item())
        gmin.append(w.min().item())
print(f"\n\nglobal max:{np.max(gmax)} min:{np.min(gmin)}")        
EOF
python3 dump_weight.py

2.检测训练过程中的异常值

A.通过hook module,检测异常值

B.拦截算子,检测异常值,打印调用栈,保存输入参数,方便复现

C.拦截算子,同时执行cpu计算,对比误差,找到第一个精度异常的算子

D.以上的代码

import torch
from torch import nn
import math
import copy
import torch.nn.functional as F

device="xpu"    
if torch.cuda.is_available():
    device="cuda"

def check_tensor(tensor, module_name, hook_type):
    if isinstance(tensor, torch.Tensor):
        if not torch.isfinite(tensor).all():
            print(f"[ERROR] Detected NaN or Inf in {hook_type} pass of {module_name} rank:{torch.distributed.get_rank()}")
            #os._exit(0)
    elif isinstance(tensor, list) or isinstance(tensor, tuple):
        for t in tensor:
            check_tensor(t, module_name, hook_type)

# 定义钩子函数来监测 NaN 和 Inf
def forward_hook(module, inputs, output, module_name):
    check_tensor(inputs, module_name, "forward-inputs")
    check_tensor(output, module_name, "forward-output")

def backward_hook(module, grad_input, module_name):
    check_tensor(grad_input, module_name, "backward")

from torch.utils._python_dispatch import TorchDispatchMode
from dataclasses import dataclass
from typing import Any
import inspect

@dataclass
class _ProfilerState:
    cls: Any
    object: Any = None

def is_valid(val,name,stack):
    # 判断是否为tensor或Parameter
    if isinstance(val, (torch.Tensor, nn.Parameter)):
        if not torch.isfinite(val.cpu()).all():
            print("[ERROR]:",name,stack)
            return -1
    return 0

def check_tensor(name,stack, tensor):
    if isinstance(tensor,(torch.Tensor, nn.Parameter)):
      return is_valid(tensor,name,stack)
    elif isinstance(tensor, (tuple, list)):
        for idx, t in enumerate(tensor):
            if is_valid(t,name,stack)!=0:
                return -1
    return 0

def save_tensor_data(tensor,name):
    if isinstance(tensor,(torch.Tensor, nn.Parameter)):
	    torch.save(tensor,f"{name}.pth")
    elif isinstance(tensor, (tuple, list)):
        for idx, t in enumerate(tensor):
            save_tensor_data(t,f"{name}-{idx}")

def to_cpu_and_fp32(data):
    """
    将输入数据中的所有GPU Tensor转换为CPU Tensor，并将FP16类型的张量转换为FP32类型。

    参数：
    data: 可能是单个Tensor，列表，元组，字典或这些数据类型的嵌套结构。

    返回：
    与输入结构相似，但所有的GPU张量都已转换为CPU张量，所有的FP16张量都转换为FP32张量。
    """
    if isinstance(data, torch.Tensor):
        # 将GPU张量转换为CPU张量
        tensor = data.cpu()
        # 如果张量是FP16类型，则转换为FP32
        if tensor.dtype == torch.float16:
            tensor = tensor.to(torch.float32)
        return tensor
    elif isinstance(data, list):
        return [to_cpu_and_fp32(item) for item in data]
    elif isinstance(data, tuple):
        return tuple(to_cpu_and_fp32(item) for item in data)
    elif isinstance(data, dict):
        return {key: to_cpu_and_fp32(value) for key, value in data.items()}
    else:
        # 如果既不是Tensor也不是列表、元组或字典，则直接返回数据
        return data

@dataclass
class TensorDesc:
    cat: Any
    shape: Any
    dtype: Any
    value: Any

@dataclass
class DataDescriptor:
    class_name: Any
    shape: Any
    value: Any
    dtype: Any
    max_v: Any
    min_v: Any    
    def __repr__(self) -> str:
        output_str=[]
        if self.shape:
            output_str.append("shape:({})".format(",".join([str(x) for x in self.shape])))
        if self.max_v:
            output_str.append(f"max:{self.max_v:.6f} min:{self.min_v:.6f}")
        if self.value is not None:
            if self.class_name in ["list","tuple"]:
                for t in self.value:
                    output_str.append(str(t))
            else:
                output_str.append(str(self.value))
        if self.dtype and self.class_name in ["Tensor","ndarray","Parameter"]:
            output_str.append(str(self.dtype))
        return "{}({})".format(self.class_name,"-".join(output_str))

class InputDescriptor:
    def __init__(self) -> None:
        self.input_vars=[]
        self.input_kwargs={}
        
    def _save_var(self,v):
        class_name=v.__class__.__name__
        if class_name in ["Tensor","Parameter"]:
            return DataDescriptor(class_name,list(v.shape),None,v.dtype,v.max().item(),v.min().item())
        elif class_name in ["UntypedStorage"]:
            pass
            #return DataDescriptor(class_name,None,list(v),type(v))        
        elif class_name in ["int","float","str","dtype","layout","device","NoneType","bool","memory_format"]:
            return DataDescriptor(class_name,None,v,type(v),None,None)
        elif class_name in ["ndarray"]:
            return DataDescriptor(class_name,list(v.shape),None,v.dtype,None,None)
        elif class_name in ["list","tuple"]:
            output=[]
            for t in v:
                output.append(self._save_var(t))
            return DataDescriptor(class_name,None,output,None,None,None)
        
    def save_vars(self,*args,**kwargs):        
        for arg in args:
            self.input_vars.append(self._save_var(arg))
        for k,v in kwargs.items():
            self.input_kwargs[k]=self._save_var(v)

    def __repr__(self) -> str:
        return str(self.input_vars) + "#" + str(self.input_kwargs)

def compare_tensor(tensorA,tensorB,name,params):
    if isinstance(tensorA,(torch.Tensor, nn.Parameter)):
        mse_loss = torch.nn.MSELoss()
        loss = mse_loss(tensorA.cpu().float(), tensorB.cpu().float()).item()
        print(f"{name:<64s} {loss:f} {params}")
        return loss<1e-2
    elif isinstance(tensorA, (tuple, list)):
        for idx, t in enumerate(tensorA):
            if not compare_tensor(tensorA[idx],tensorB[idx],f"{name}-{idx}",params):
                return False
    return True
    
def is_in_blacklist(name):
    black_list=["empty","like","zero","detach","has","view",
                "copy","arange","fill","ones","lift_fresh","alias",
                "scalar_tensor","clone","stack","slice","source",
                "select","random","unsqueeze","expand","normal","bernoulli"]
    for i in black_list:
        if name.find(i)>=0:
            return False
    return True
    
class TorchOpDiffDispatchMode(TorchDispatchMode):
    def __init__(self,parent):
        super().__init__()
        self.parent=parent
        self.global_index=0

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        func_packet = func._overloadpacket
        op_name=f"{func}"
        enable_dump= is_in_blacklist(op_name)
        self.global_index+=1
        if kwargs is None:
            kwargs = {}
        if enable_dump:
            args_cpu = to_cpu_and_fp32(args)
            kwargs_cpu = to_cpu_and_fp32(kwargs)
            cpu_out=func(*args_cpu, **kwargs_cpu)
        ret= func(*args, **kwargs)
        if enable_dump:
            desc=InputDescriptor()
            desc.save_vars(*args,**kwargs)        
            if not compare_tensor(cpu_out,ret,op_name,str(desc)):
                save_tensor_data(args,f"{self.global_index}_{torch.distributed.get_rank()}{op_name}-input")
                save_tensor_data(ret,f"{self.global_index}_{torch.distributed.get_rank()}{op_name}-output")
        return ret

class TorchNanDetDispatchMode(TorchDispatchMode):
    def __init__(self,parent):
        super().__init__()
        self.parent=parent
        self.global_index=0

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        func_packet = func._overloadpacket   
        op_name=f"{func}"  
        self.global_index+=1
        enable_dump= is_in_blacklist(op_name)
        if kwargs is None:
            kwargs = {}
        stacks=[i for i in inspect.stack()]
        stacks_sz=len(stacks)
        msg=[]
        for idx in range(stacks_sz-1,1,-1):
            if "self" in stacks[idx].frame.f_locals:
                class_name = stacks[idx].frame.f_locals["self"].__class__.__name__
            else:
                class_name=""
            msg.append(f"{stacks[idx].filename}:[{class_name}]:{stacks[idx].function}")        
        valid=0
        if enable_dump:
            valid+=check_tensor(f"aten-{op_name}-input","\n".join(msg),args)    
        ret= func(*args, **kwargs)
        if enable_dump:
            valid+=check_tensor(f"{op_name}-output","\n".join(msg), ret)
        if valid!=0:
            save_tensor_data(args,f"{self.global_index}_{torch.distributed.get_rank()}{op_name}-input")
            save_tensor_data(ret,f"{self.global_index}_{torch.distributed.get_rank()}{op_name}-output")
        return ret

class TorchHook:
    _CURRENT_Dumper = None
    def __init__(self,state):
        self.p= _ProfilerState(state)
        
    def __enter__(self):
        assert TorchHook._CURRENT_Dumper is None
        TorchHook._CURRENT_Dumper = self
        if self.p.object is None:
            o = self.p.cls(self)
            o.__enter__()
            self.p.object = o
        else:
            self.p.object.step()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        TorchHook._CURRENT_Dumper = None
        if self.p.object is not None:
            self.p.object.__exit__(exc_type, exc_val, exc_tb)
            del self.p.object

def clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self,query, key, value, mask=None, dropout=None):
        d_k = query.size(-1)
        scores = query@key.transpose(-2,-1) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e20)
        p_attn = F.softmax(scores, dim = -1)
        if dropout is not None:
            p_attn = dropout(p_attn)
        return p_attn@value, p_attn

class MultiHeadAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
        self.attention = ScaledDotProductAttention()

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        query=self.linears[0](query).view(nbatches, -1, self.h, self.d_k)
        query=query.transpose(1, 2)
        key=self.linears[1](key).view(nbatches, -1, self.h, self.d_k)
        key=key.transpose(1, 2)
        value=self.linears[2](value).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
        x, self.attn = self.attention(query, key, value, mask=mask,
                                 dropout=self.dropout)
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

torch.random.manual_seed(1)
model = MultiHeadAttention(h=8, d_model=64).half().to(device)
model.eval()

# 通过hook Module去检测是否有Nan值
for name, module in model.named_modules():
    if isinstance(module, nn.Module):  # 检测是否是 nn.Module 的子类
        module.register_forward_hook(lambda module, inputs, output, name=name: forward_hook(module, inputs, output, name))
        module.register_full_backward_pre_hook(lambda module, grad_input, name=name: backward_hook(module, grad_input, f"{name}-in"))
        module.register_full_backward_hook(lambda module, grad_input,name=name: backward_hook(module, grad_input, f"{name}-out"))

q1 = torch.ones((100, 50, 64),dtype=torch.float32).half().to(device)
k1 = q1.clone()
v1 = q1.clone()

# 通过拦截算子去检测是否有nan值,并打印调用栈,保存输入参数,方便后续复现
with TorchHook(TorchNanDetDispatchMode):
    out = model(q1,k1,v1).sum()
    print("out:",out.item())

# 通过与cpu比较算子的计算误差,保存输入参数
with TorchHook(TorchOpDiffDispatchMode):
    out = model(q1,k1,v1).sum()
    print("out:",out.item())

3.根据上面dump的数据,准备最小复现环境

tee demo.py <<-'EOF'
import torch
import numpy as np

device="xpu"

def largest_k_errors(tensor1, tensor2, k, eps=1e-12):
    """
    找到两个张量之间误差最大的 k 个元素及其原始数据。
    
    参数：
    tensor1: 第一个输入张量。
    tensor2: 第二个输入张量。
    k (int): 要找的误差最大的元素个数。
    eps (float): 防止除以零的一个很小的值，默认值为1e-12。
    
    返回：
    包含误差最大的 k 个元素及其原始数据的元组 (errors, indices, orig_values1, orig_values2)。
    """
    # 对输入张量进行广播操作，使它们的形状一致
    tensor1, tensor2 = torch.broadcast_tensors(tensor1, tensor2)

    # 计算绝对误差
    absolute_diff = torch.abs(tensor1 - tensor2)
    
    # 计算基准值，避免除以零的情况
    denom = torch.abs(tensor1) + torch.abs(tensor2) + eps
    
    # 计算相对误差
    relative_error = absolute_diff / denom
    
    # 把inf和nan替换为零
    relative_error = torch.where(
        torch.isfinite(relative_error),
        relative_error,
        torch.zeros_like(relative_error)
    )
    
    # 找到误差最大的 k 个元素的索引
    k = min(k, tensor1.numel())  # 确保 k 不超过张量的元素总数
    _, indices = torch.topk(relative_error.view(-1), k)
    
    # 提取原始数据和误差
    errors = relative_error.view(-1)[indices]
    orig_values1 = tensor1.view(-1)[indices]
    orig_values2 = tensor2.view(-1)[indices]

    return errors, indices, orig_values1, orig_values2

left=torch.load("8611_0aten.silu_backward.default-input-0.pth",map_location="cpu")
right=torch.load("8611_0aten.silu_backward.default-input-1.pth",map_location="cpu")
tensor1=torch.ops.aten.silu_backward.default(left,right).cpu().float().reshape(-1)
out_xpu=torch.ops.aten.silu_backward.default(left.to(device),right.to(device))
tensor2=out_xpu.cpu().float().reshape(-1)

k = 10  # 找出误差最大的K个元素

errors, indices, orig_values1, orig_values2 = largest_k_errors(tensor1, tensor2, k)

print("----------------- torch.ops.aten.silu_backward.default --------------------------")
print(f"left  dtype:{left.dtype} stride:{left.stride} shape:{left.shape} max:{left.max().item()} min:{left.min().item()}")
print(f"right dtype:{right.dtype} shape:{right.shape} max:{right.max().item()} min:{right.min().item()}")

print("Top", k, "errors:")
for i in range(k):
    print(f"Index {indices[i].item()}: Error {errors[i].item()}, "
        f"Original Values: CPU_FP16 = {orig_values1[i].item()}, "
        f"XPU_FP16 = {orig_values2[i].item()}")
EOF
python3 demo.py