如何让PyTorch模型返回float32/float16张量，解决LLM微调验证内存不足

• 加载模型时指定dtype
  直接在from_pretrained中设置torch_dtype，让模型参数和输出默认使用目标精度，这是最直接的方式：

  import torch
  from transformers import AutoModelForCausalLM
  
  # 加载为float32
  model = AutoModelForCausalLM.from_pretrained(
      "your-model-path",
      torch_dtype=torch.float32
  )
  
  # 加载为float16（GPU场景更推荐，内存占用减半）
  model = AutoModelForCausalLM.from_pretrained(
      "your-model-path",
      torch_dtype=torch.float16
  ).to("cuda")
  

• 验证循环中强制转换输出张量
  如果不想修改模型整体dtype，可以在验证阶段单独转换模型输出的张量类型：

  for batch in val_dataloader:
      batch = {k: v.to(device) for k, v in batch.items()}
      with torch.no_grad():
          outputs = model(**batch)
          # 将logits转换为float32/float16
          logits = outputs.logits.to(torch.float32)  # 替换为torch.float16即可
          # 后续损失计算、指标评估都基于转换后的张量
          loss = loss_function(logits, batch["labels"])
  

• 调整损失函数的精度
  部分损失函数默认会使用float64计算，手动指定其dtype可以避免额外的内存开销：

  from torch.nn import CrossEntropyLoss
  
  # 让损失函数使用float32计算
  loss_function = CrossEntropyLoss().to(torch.float32)
  

• 全局设置默认浮点类型（慎用）
  可以通过PyTorch全局设置改变默认的浮点 dtype，但可能影响其他未适配的代码，仅推荐在全流程统一精度的场景使用：

  import torch
  # 设置默认浮点类型为float32
  torch.set_default_dtype(torch.float32)
  

注意事项：

• float16精度仅建议在GPU上使用，CPU对float16的支持有限，可能导致性能下降或报错；
• float16存在一定精度损失，但在大多数LLM微调任务中，这种损失不会显著影响最终效果；
• 使用float16时，建议配合torch.cuda.amp.autocast()上下文管理器，确保部分算子的数值稳定性：

  with torch.no_grad(), torch.cuda.amp.autocast():
      outputs = model(**batch)
  

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