Apache Spark 3.4 为深度学习（DL）工作负载带来了原生支持，让大规模数据上的分布式模型训练与推理变得前所未有的简便。本教程将深入介绍两个核心 API：TorchDistributor（用于分布式训练）和 predict_batch_udf（用于分布式推理），并通过代码示例和常见问题解答帮助你快速上手。

分布式训练：TorchDistributor

在 Spark 3.4 之前，将 PyTorch 分布式训练代码迁移到 Spark 需要依赖第三方库（如 Horovod-on-Spark）。如今，Spark 内置了 TorchDistributor API，它借助 Spark 的屏障执行模式，在 Spark Executor 上自动启动分布式深度学习集群节点。你只需定义主训练函数 main_fn，然后调用 TorchDistributor.run() 即可。

核心优势：现有的分布式训练代码几乎无需修改，即可在 Spark 上运行。训练进程之间能够互相通信，并直接访问 Spark 集群的分布式文件系统。

from pyspark.ml.torch.distributor import TorchDistributor

def main_fn(checkpoint_dir):
    # standard distributed PyTorch code
    ...

# Set num_processes = NUM_WORKERS * NUM_GPUS_PER_WORKER
output_dist = TorchDistributor(
    num_processes=2, 
    local_mode=False, 
    use_gpu=True
).run(main_fn, checkpoint_dir)

图 1：使用 TorchDistributor API 进行分布式训练

迁移指南与注意点

• 代码修改量最少：你只需将原本的单机分布式代码放入 main_fn 中，并调整读取数据的路径为分布式文件系统路径。
• 数据处理提前完成：TorchDistributor 不使用 Spark RDD 或 DataFrame 传输数据，因此所有 Spark 预处理必须在训练启动前完成并持久化到存储（如 HDFS、S3）。
• 通信依赖框架：Executor 上的进程使用 PyTorch 内置的分布式通信（如 NCCL、GLOO）进行梯度交换。