使用Accelerate FSDP时单进程CLIPVisionModel权重加载错误致训练NaN

在2张48G显存A40 GPU上使用Accelerate FSDP训练时，出现训练损失为NaN、权重为-inf的情况。调试发现初始化阶段加载的非训练模型CLIPVisionModel在其中一个进程中权重异常，进而导致隐藏参数出现inf/NaN，最终引发训练损失NaN。

权重对比 #

正常进程CLIPVisionModel权重示例：

tensor([0.3311, 0.0032, 0.1610, ..., 2.1922, 0.0050, 0.0039],)

异常进程CLIPVisionModel权重示例：

tensor([-1.9921e-04, 4.5673e-41, -1.9921e-04, ..., 0.0000e+00,0.0000e+00, 0.0000e+00],)

加载代码 #

class SLlamaModel(LlamaModel):
    config_class = SConfig

    def __init__(self, config: LlamaConfig, mm_vision_tower=None, mm_hidden_size=None):
        super(SLlamaModel, self).__init__(config)

        if hasattr(config, "mm_vision_tower"):
            # HACK: for FSDP
            self.vision_tower = CLIPVisionModel.from_pretrained(config.mm_vision_tower)

        if hasattr(config, "use_mm_proj"):
            self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)

注：config.mm_vision_tower为clip-vit-large-patch14的本地路径

FSDP配置 #

• compute_environment: LOCAL_MACHINE
• distributed_type: FSDP
• downcast_bf16: 'no'
• fsdp_config:
  • fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  • fsdp_backward_prefetch_policy: BACKWARD_PRE
  • fsdp_forward_prefetch: true
  • fsdp_offload_params: true
  • fsdp_sharding_strategy: 1
  • fsdp_state_dict_type: FULL_STATE_DICT
  • fsdp_sync_module_states: true
  • fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  • fsdp_use_orig_params: true
• machine_rank: 0
• main_training_function: main
• mixed_precision: bf16
• num_machines: 1
• num_processes: 2
• rdzv_backend: static
• same_network: true
• tpu_env: []
• tpu_use_cluster: false
• tpu_use_sudo: false
• use_cpu: false

环境信息 #

• Accelerate version: 0.21.0
• Platform: Linux-5.4.0-90-generic-x86_64-with-glibc2.31
• Python version: 3.10.12
• Numpy version: 1.25.1
• PyTorch version (GPU?): 2.0.1+cu117 (False)
• PyTorch XPU available: False
• PyTorch NPU available: False
• System RAM: 755.74 GB

调试结论 #

经全流程调试确认，权重错误源于from_pretrained加载阶段，而非后续代码问题——直接调用test = CLIPVisionModel.from_pretrained("model path")仍会得到异常权重。

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