PyTorch+FFCV训练中_thread.lock.acquire过慢致迭代时长波动问题

问题背景 #

使用Python、PyTorch及FFCV数据加载器在单块NVIDIA GPU上训练神经网络时，发现单次迭代时长差异可达30倍，同一训练任务内也会出现时快时慢的情况，导致整体训练耗时大幅增加。通过cProfile性能分析发现，耗时最高的函数为method 'acquire' of '_thread.lock' objects，单次调用耗时从快速轮次的0.015s到慢速轮次的0.124s，差异约10倍。

性能分析结果 #

小型训练（快速） #

6618917 function calls (5423678 primitive calls) in 119.280 seconds

   Ordered by: cumulative time
   List reduced from 1214 to 100 due to restriction <100>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.065    0.065  119.280  119.280 ./train.py:261(do_training)
     1565    0.018    0.000   92.495    0.059 .../.conda/envs/ffcv/lib/python3.9/site-packages/ffcv/loader/epoch_iterator.py:140(__next__)
     1565    0.023    0.000   92.258    0.059 .../.conda/envs/ffcv/lib/python3.9/queue.py:154(get)
     1142    0.015    0.000   92.211    0.081 .../.conda/envs/ffcv/lib/python3.9/threading.py:280(wait)
     6134   92.190    0.015   92.190    0.015 {method 'acquire' of '_thread.lock' objects}
     1564    0.134    0.000   22.608    0.014 ./models/ssl.py:560(step)
136068/3128    0.468    0.000   21.914    0.007 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/module.py:1124(_call_impl)
     1562    0.022    0.000   20.810    0.013 ./models/ssl.py:259(test_step)
     1564    0.279    0.000   15.438    0.010 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py:147(forward)
     1564    0.010    0.000   11.614    0.007 ./models/encoder.py:55(forward)
     1564    0.074    0.000   11.593    0.007 ./models/resnet.py:32(forward)
10948/6256    0.076    0.000    8.878    0.001 ...e/.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/container.py:137(forward)
    12512    0.529    0.000    8.759    0.001 .../.conda/envs/ffcv/lib/python3.9/site-packages/torchvision/models/resnet.py:89(forward)
    31280    0.095    0.000    7.447    0.000 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/conv.py:456(forward)
    31280    0.078    0.000    7.315    0.000 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/conv.py:448(_conv_forward)
    31280    7.237    0.000    7.237    0.000 {built-in method torch.conv2d}

小型训练（慢速） #

2570224 function calls (2396817 primitive calls) in 1005.364 seconds

   Ordered by: cumulative time
   List reduced from 1250 to 100 due to restriction <100>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.076    0.076 1005.364 1005.364 ./train.py:261(do_training)
     1565    0.025    0.000  972.119    0.621 .../.conda/envs/ffcv/lib/python3.9/site-packages/ffcv/loader/epoch_iterator.py:140(__next__)
     1565    0.033    0.000  971.742    0.621 .../.conda/envs/ffcv/lib/python3.9/queue.py:154(get)
     1565    0.026    0.000  971.675    0.621 .../.conda/envs/ffcv/lib/python3.9/threading.py:280(wait)
     7827  971.640    0.124  971.640    0.124 {method 'acquire' of '_thread.lock' objects}
     1564    0.159    0.000   26.123    0.017 ./models/ssl.py:560(step)
136068/3128    0.650    0.000   25.136    0.008 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/module.py:1124(_call_impl)
     1562    0.021    0.000   22.797    0.015 ./models/ssl.py:259(test_step)
     1564    0.104    0.000   18.512    0.012 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/parallel/distributed.py:973(forward)
     1564    0.023    0.000   17.731    0.011 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/parallel/distributed.py:959(_run_ddp_forward)
     1564    0.014    0.000   17.501    0.011 ./models/encoder.py:55(forward)
     1564    0.120    0.000   17.466    0.011 ./models/resnet.py:32(forward)
10948/6256    0.111    0.000   14.324    0.002 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/container.py:137(forward)
    12512    0.756    0.000   14.152    0.001 .../.conda/envs/ffcv/lib/python3.9/site-packages/torchvision/models/resnet.py:89(forward)
    31280    0.142    0.000   11.407    0.000 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/conv.py:456(forward)
    31280    0.107    0.000   11.207    0.000 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/nn/modules/conv.py:448(_conv_forward)
    31280   11.100    0.000   11.100    0.000 {built-in method torch.conv2d}

大型训练（慢速） #

9393663 function calls (9055831 primitive calls) in 2093.533 seconds

   Ordered by: cumulative time
   List reduced from 1276 to 100 due to restriction <100>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      335    0.007    0.000 1942.907    5.800 .../.conda/envs/ffcv/lib/python3.9/threading.py:280(wait)
      344    0.007    0.000 1942.901    5.648 .../.conda/envs/ffcv/lib/python3.9/site-packages/ffcv/loader/epoch_iterator.py:140(__next__)
     1692 1942.898    1.148 1942.898    1.148 {method 'acquire' of '_thread.lock' objects}
      344    0.008    0.000 1942.836    5.648 .../.conda/envs/ffcv/lib/python3.9/queue.py:154(get)
       42    0.002    0.000   48.570    1.156 ./train.py:238(save_checkpoint)
       84    0.008    0.000   48.519    0.578 ./train.py:362(save_checkpoint_data)
       84    0.016    0.000   48.490    0.577 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/serialization.py:333(save)
      160    0.001    0.000   44.991    0.281 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/_tensor.py:340(backward)
      160    0.001    0.000   44.990    0.281 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/autograd/__init__.py:85(backward)
      160   44.985    0.281   44.985    0.281 ...._EngineBase' objects}
       84    0.001    0.000   41.014    0.488 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/serialization.py:213(__exit__)
       84   41.014    0.488   41.014    0.488 {method 'close' of '_io.BufferedWriter' objects}
      160    0.001    0.000   20.002    0.125 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py:288(step)
      160    0.002    0.000   20.000    0.125 .../.conda/envs/ffcv/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py:282(_maybe_opt_step)
     1953    0.008    0.000   18.784    0.010 {built-in method builtins.sum}

成因分析 #

• FFCV多线程数据加载的锁竞争：FFCV依赖多线程进行数据预处理与加载，主线程通过队列获取预处理好的batch时需要获取锁。当数据生产速度跟不上GPU的消费速度（比如GPU训练效率高、预处理线程被抢占），主线程会陷入长时间等待锁的状态，导致acquire耗时剧增。
• 系统资源调度干扰：训练过程中如果有其他进程抢占CPU、内存或磁盘IO资源，会导致FFCV的预处理线程无法及时完成数据处理，队列空转，主线程等待锁的时间被拉长。
• 训练流程中的同步操作：比如检查点保存（从大型训练的分析结果可见，save_checkpoint占用了不少时间），这类操作会占用大量IO和CPU资源，间接影响数据加载线程的运行效率，加剧锁等待。
• PyTorch异步操作的波动：CUDA流的同步、梯度计算的延迟等情况，会让GPU的空闲时间变长，主线程等待数据的时间也随之增加，锁的acquire操作耗时自然上升。

缓解与解决方法 #

优化FFCV数据加载器配置 #

• 增大num_workers：增加预处理线程数量，确保数据生产速度能匹配GPU的消费速度，减少主线程等待队列的时间。
• 提高prefetch_factor：调整队列预取的batch数量，让队列中始终保持足够的待处理batch，避免主线程空等。
• 启用pin_memory：减少CPU到GPU的数据拷贝耗时，让GPU更快完成计算，间接降低数据加载的压力。

系统资源优化 #

• 绑定CPU核心：用taskset工具将训练进程绑定到特定CPU核心，避免其他进程抢占CPU资源，保证预处理线程的运行稳定性。
• 升级存储设备：将训练数据迁移到SSD或内存盘（tmpfs），降低磁盘IO延迟，提升数据读取速度。
• 清理后台进程：关闭不必要的后台服务，释放CPU、内存和IO资源，减少对训练进程的干扰。

训练流程调整 #

• 调整检查点保存策略：降低检查点保存频率，或者用单独线程异步保存检查点，避免阻塞训练主线程。
• 减少不必要的同步操作：移除训练代码中多余的torch.cuda.synchronize()调用，让GPU与CPU操作异步执行，提升整体效率。

FFCV预处理管道优化 #

• 简化预处理步骤：移除不必要的数据预处理操作，将可并行的预处理任务分配给更多线程。
• 优化数据编码：确保FFCV数据集使用最优的编码格式（如合适的JPEG质量），减少数据解码耗时。

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