Datacrunch上TensorFlow分布式训练GPU利用率低且不稳定排查

从锯齿状大幅波动的GPU利用率曲线可以判断，核心问题是数据流水线供数速度远低于GPU计算速度，GPU频繁处于“算完一批等待下一批”的空闲状态，具体诱因如下：

• tf.data流水线未开启并行处理：当前ds.map()调用未指定num_parallel_calls参数，默认单线程执行批次解析、数据增强逻辑，CPU处理速度完全跟不上两块A100的消费速度。
• 数据增强逻辑存在严重性能缺陷：augment函数中使用Python原生random库生成随机数，且未做静态图编译，这部分逻辑无法被TensorFlow纳入高效图执行流程，会频繁触发Python运行时和TF运行时的上下文切换，且不支持多线程并行，单步处理延迟极高。如果_reverb混响函数基于Python/NumPy实现，还会额外增加大量CPU侧串行计算开销。
• 分布式训练未做数据集分片：使用MirroredStrategy时直接传入普通数据集，会导致两块GPU重复拉取全量数据，CPU侧和IO侧的供数压力翻倍，进一步拉大供数和计算的速度差。
• 流水线参数配置不合理：未开启数据预取到GPU的优化，IO、数据预处理、GPU计算三个环节完全串行，没有重叠执行，空等时间长。

按优先级从高到低落地以下调整，GPU利用率可稳定到85%以上：

1. 重构tf.data流水线，拉满并行能力，适配分布式场景
   调整流水线执行顺序，开启全链路并行，按卡分片数据集，让IO、预处理、计算完全重叠，参考代码如下：

   # 开启tf.data自动调优
   options = tf.data.Options()
   options.autotune.enabled = True
   options.deterministic = False  # 训练阶段无需严格保证样本顺序，换取更高吞吐量
   
   global_batch_size = ... # 原有全局批次大小配置
   def build_distributed_dataset(input_context):
       # 每个GPU分走独立的文件分片，避免重复读取数据
       per_replica_batch = input_context.get_per_replica_batch_size(global_batch_size)
       sharded_files = files_ds.shard(
           num_shards=input_context.num_input_pipelines,
           index=input_context.input_pipeline_id
       )
       ds = tf.data.TFRecordDataset(
           sharded_files,
           compression_type='ZLIB',
           num_parallel_reads=tf.data.AUTOTUNE
       )
       # 加样本混洗，缓冲区大小设为单TFRecord样本量的4倍即可
       ds = ds.shuffle(buffer_size=2400)
       ds = ds.batch(per_replica_batch, drop_remainder=True)
       # 开启map多线程并行
       ds = ds.map(
           lambda x: _parse_batch(x, sample_rate, duration, split),
           num_parallel_calls=tf.data.AUTOTUNE,
           deterministic=False
       )
       if args.steps_per_epoch != -1:
           ds = ds.repeat()
       # 预取数据直接放到对应GPU显存，减少PCIe传输等待
       ds = ds.prefetch(buffer_size=tf.data.AUTOTUNE)
       ds = ds.apply(tf.data.experimental.prefetch_to_device(f"/gpu:{input_context.input_pipeline_id}"))
       return ds.with_options(options)
   
   # 用分布式接口包装数据集，自动为每个卡构建独立流水线
   ds = strategy.distribute_datasets_from_function(build_distributed_dataset)
   

2. 重构数据增强逻辑，全量替换为TF原生算子
   删掉所有Python原生随机数、NumPy计算逻辑，全部换成TF原生API，加编译装饰器转为静态图执行，让这部分逻辑支持多线程并行：

   @tf.function(jit_compile=True) # 开启XLA编译，进一步提速
   def augment(noisy, clean):
       proba_threshold = 0.5
       initial_max = 0.3
       rt60_range = (0.3, 1.3)
       first_delay_range = (0.01, 0.03)
       repeat_times = 3
       jitter = 0.1
       keep_clean_ratio = 0.1
       sr = 48000
   
       # 替换Python random为TF原生随机算子
       if tf.random.uniform(()) >= proba_threshold:
           return noisy, clean
       noise = noisy - clean
       initial = tf.random.uniform(()) * initial_max
       first_delay = tf.random.uniform((), first_delay_range[0], first_delay_range[1])
       rt60 = tf.random.uniform((), rt60_range[0], rt60_range[1])
       # 注意_reverb函数也要全部用TF原生算子实现，不能留Python/NumPy逻辑
       reverb_noise = _reverb(noise, initial, first_delay, rt60, repeat_times, jitter, sr)
       noise += reverb_noise
       reverb_clean = _reverb(clean, initial, first_delay, rt60, repeat_times, jitter, sr)
       clean += keep_clean_ratio * reverb_clean
       noise += (1 - keep_clean_ratio) * reverb_clean
       noisy = noise + clean
       return noisy, clean
   

3. 把重计算逻辑挪到GPU侧执行
   混响这类信号处理计算本质是矩阵运算，CPU执行效率极低，可以把augment逻辑从数据预处理的map环节拆出来，作为模型的第一层放到前向传播逻辑里，直接在GPU上执行，CPU侧只保留TFRecord解压、样本解析这类轻量操作。
4. 额外性能调优
   • 开启A100适配的bfloat16混合精度训练，单步训练速度可以提升40%以上，GPU利用率更稳定：

     tf.keras.mixed_precision.set_global_policy('mixed_bfloat16')
     

   • 适当调大全局batch size：A100单卡显存可达40G/80G，处理10s 48k音频时单卡batch size可以开到32~64，避免batch太小导致GPU计算速度过快、频繁等数据。
   • 如果CPU核数充足，可以把TFRecord的压缩格式从ZLIB换成LZ4，解压速度提升3~5倍，大幅降低IO解压开销。

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