分布式训练TorchDistributor
本文描述如何执行分布式训练PyTorch毫升模型使用TorchDistributor。
TorchDistributor PySpark开源模块,帮助用户做分布式训练PyTorch火花集群,这允许您启动PyTorch培训工作引发的工作。底层,它初始化环境和员工之间的沟通渠道,充分利用了CLI命令torch.distributed.run工作者节点运行分布式训练。
的TorchDistributor API支持方法如下表所示。
方法和签名 |
描述 |
|---|---|
|
TorchDistributor创建一个实例。 |
|
通过调用运行分布式训练 |
开发工作流程为笔记本电脑
如果模型创建和训练过程完全从你的本地机器或笔记本砖的笔记本,你只需要准备你的代码进行少量更改为分布式训练。
准备单独的节点代码:准备和测试单一节点代码PyTorch, PyTorch闪电,或其他框架基于PyTorch / PyTorch闪电,HuggingFace教练API。
准备代码标准分布式训练:你需要单流程培训转化成分布式训练。这种分布式代码都包含在一个训练函数,您可以使用
TorchDistributor。移动进口在训练函数:添加必要的进口,如
进口火炬在培训的功能。这样做可以让你避免常见的酸洗错误。此外,device_id模型和数据的联系是由:device_id=int(操作系统。环境(“LOCAL_RANK”])
启动分布式训练:实例化
TorchDistributor所需的参数和调用.run (* args)启动培训。
以下是培训代码示例:
从pyspark.ml.torch.distributor进口TorchDistributordef火车(learning_rate,use_gpu):进口火炬进口torch.distributed作为经销进口torch.nn.parallel.DistributedDataParallel作为DDP从torch.utils.data进口DistributedSampler,DataLoader后端=“nccl”如果use_gpu其他的“gloo”经销。init_process_group(后端)设备=int(操作系统。环境(“LOCAL_RANK”])如果use_gpu其他的“cpu”模型=DDP(createModel(),* *kwargs)取样器=DistributedSampler(数据集)加载程序=DataLoader(数据集,取样器=取样器)输出=火车(模型,加载程序,learning_rate)经销。清理()返回输出经销商=TorchDistributor(num_processes=2,local_mode=假,use_gpu=真正的)经销商。运行(火车,1 e - 3,真正的)
从外部存储库迁移训练
如果你有一个现有的分布式训练程序存储在外部存储库中,您可以很容易地迁移到砖做的以下几点:
导入库:导入外部存储库是一个砖回购。
创建一个新的笔记本初始化一个新的库中的数据砖的笔记本。
启动分布式训练笔记本电池,电话
TorchDistributor如下:
从pyspark.ml.torch.distributor进口TorchDistributortrain_file=“/道路/ / train.py”arg游戏=(“——learning_rate = 0.001”,”——batch_size = 16 "]经销商=TorchDistributor(num_processes=2,local_mode=假,use_gpu=真正的)经销商。运行(train_file,*arg游戏)