炼丹师的优雅内功二：PyTorch on kubernetes

前面的文章

炼丹师的优雅内功：从机器学习，分布式训练到k8s云计算

炼丹师的优雅内功一：tensorflow on kubernetes

一、PyTorch 基础

1.1 是什么

PyTorch是一个开源的Python机器学习库，基于Torch，底层由C++实现，应用于人工智能领域，如自然语言处理。 它最初由Facebook的人工智能研究团队开发，并且被用于Uber的概率编程软件Pyro。 PyTorch主要有两大特征： 类似于NumPy的张量计算，可使用GPU加速.

1.2 基础知识

1.2.1 张量(Tensor)

PyTorch 中的张量和tensorflow 含义相同，都是n 维数组，pytorch 中的张量计算可以通过GPU 加速计算 张量维度

维度	含义
0	代表标量(数字)
1	代表向量()
2	代表矩阵
3	时间序列数据，例如股价 文本数据 单张彩色图片
4	图像
5	视频
1.2.1.1 创建tersor

• 1. 随机初始化矩阵 我们可以通过torch.rand()的方法，构造一个随机初始化的矩阵：

import torch
x = torch.rand(4, 3) 
print(x)

### 输出
tensor([[0.7569, 0.4281, 0.4722],
        [0.9513, 0.5168, 0.1659],
        [0.4493, 0.2846, 0.4363],
        [0.5043, 0.9637, 0.1469]])

• 2. 全0矩阵的构建 我们可以通过torch.zeros()构造一个矩阵全为 0，并且通过dtype设置数据类型为 long。除此以外，我们还可以通过torch.zero_()和torch.zeros_like()将现有矩阵转换为全0矩阵.

import torch
x = torch.zeros(4, 3, dtype=torch.long)
print(x)

### 输出
tensor([[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]])

• 3. 张量的构建 我们可以通过torch.tensor()直接使用数据，构造一个张量：

import torch
x = torch.tensor([5.5, 3]) 
print(x)
tensor([5.5000, 3.0000])

• 4. 基于已经存在的 tensor，创建一个 tensor

x = x.new_ones(4, 3, dtype=torch.double) 
# 创建一个新的全1矩阵tensor，返回的tensor默认具有相同的torch.dtype和torch.device
# 也可以像之前的写法 x = torch.ones(4, 3, dtype=torch.double)
print(x)
x = torch.randn_like(x, dtype=torch.float)
# 重置数据类型
print(x)
# 结果会有一样的size
# 获取它的维度信息
print(x.size())
print(x.shape)

### 输出
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]], dtype=torch.float64)
tensor([[ 2.7311, -0.0720,  0.2497],
        [-2.3141,  0.0666, -0.5934],
        [ 1.5253,  1.0336,  1.3859],
        [ 1.3806, -0.6965, -1.2255]])
torch.Size([4, 3])
torch.Size([4, 3])

1.2.1.2 张量操作

• 加法

import torch
# 方式1
y = torch.rand(4, 3) 
print(x + y)

# 方式2
print(torch.add(x, y))

# 方式3 in-place，原值修改
y.add_(x) 
print(y)

• 维度变换

x = torch.randn(4, 4)
y = x.view(16)
z = x.view(-1, 8) # -1是指这一维的维数由其他维度决定
print(x.size(), y.size(), z.size())
### 
torch.Size([4, 4]) torch.Size([16]) torch.Size([2, 8])

• pytorch 对张量进行操作的方法非常多后续可以查阅官方文档。。。

1.2.2 自动求导

PyTorch 中，所有神经网络的核心是 autograd 包。autograd包为张量上的所有操作提供了自动求导机制。它是一个在运行时定义 ( define-by-run ）的框架，这意味着反向传播是根据代码如何运行来决定的，并且每次迭代可以是不同的

1.3 并行计算简介

深度学习的发展离不开算力的发展，GPU的出现让我们的模型可以训练的更快，更好。所以，如何充分利用GPU的性能来提高我们模型学习的效果，这一技能是我们必须要学习的。这一节，我们主要讲的就是PyTorch的并行计算。PyTorch可以在编写完模型之后，让多个GPU来参与训练，减少训练时间。你可以在命令行使用nvidia-smi命令来查看你的GPU信息和使用情况。 CUDA是NVIDIA提供的一种GPU并行计算框架。对于GPU本身的编程，使用的是CUDA语言来实现的。但是，在我们使用PyTorch编写深度学习代码时，使用的CUDA又是另一个意思。在PyTorch使用 CUDA表示要开始要求我们的模型或者数据开始使用GPU了

1.3.1 网络结构分布到不同的设备中(Network partitioning)

在刚开始做模型并行的时候，这个方案使用的比较多。其中主要的思路是，将一个模型的各个部分拆分，然后将不同的部分放入到GPU来做不同任务的计算。其架构如下：

这里遇到的问题就是，不同模型组件在不同的GPU上时，GPU之间的传输就很重要，对于GPU之间的通信是一个考验。但是GPU的通信在这种密集任务中很难办到，所以这个方式慢慢淡出了视野

1.3.2 同一层的任务分布到不同数据中(Layer-wise partitioning)

第二种方式就是，同一层的模型做一个拆分，让不同的GPU去训练同一层模型的部分任务。其架构如下

这样可以保证在不同组件之间传输的问题，但是在我们需要大量的训练，同步任务加重的情况下，会出现和第一种方式一样的问题。

1.3.3 不同的数据分布到不同的设备中，执行相同的任务(Data parallelism)

第三种方式有点不一样，它的逻辑是，我不再拆分模型，我训练的时候模型都是一整个模型。但是我将输入的数据拆分。所谓的拆分数据就是，同一个模型在不同GPU中训练一部分数据，然后再分别计算一部分数据之后，只需要将输出的数据做一个汇总，然后再反传。其架构如下

这种方式可以解决之前模式遇到的通讯问题。现在的主流方式是数据并行的方式(Data parallelism)

二、Pytorch 如何进行分布式训练

2.1 单机单卡

model = Net()
model.cuda() # 模型显示转移到CUDA上

for image,label in dataloader:
    # 图像和标签显示转移到CUDA上
    image = image.cuda() 
    label = label.cuda()

2.1 单机多卡

• DP 首先我们来看单机多卡DP，通常使用一种叫做数据并行 (Data parallelism) 的策略，即将计算任务划分成多个子任务并在多个GPU卡上同时执行这些子任务。主要使用到了nn.DataParallel函数，它的使用非常简单，一般我们只需要加几行代码即可实现

model = Net()
model.cuda() # 模型显示转移到CUDA上

if torch.cuda.device_count() > 1: # 含有多张GPU的卡
	model = nn.DataParallel(model) # 单机多卡DP训练
        
 ## 指定gpu

# 1. model = nn.DataParallel(model, device_ids=[0,1]) # 使用第0和第1张卡进行并行训练
# 2. os.environ["CUDA_VISIBLE_DEVICES"] = "1,2" # 手动指定对程序可见的gpu 设备

2.2 多机多卡

• 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包。
• 使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型 DDP(model, device_ids=device_ids)
• 使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
• 使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练

DDP 不过通过DP进行分布式多卡训练的方式容易造成负载不均衡，有可能第一块GPU显存占用更多，因为输出默认都会被gather到第一块GPU上。为此Pytorch也提供了torch.nn.parallel.DistributedDataParallel（DDP）方法来解决这个问题。

针对每个GPU，启动一个进程，然后这些进程在最开始的时候会保持一致（模型的初始化参数也一致，每个进程拥有自己的优化器），同时在更新模型的时候，梯度传播也是完全一致的，这样就可以保证任何一个GPU上面的模型参数就是完全一致的，所以这样就不会出现DataParallel那样显存不均衡的问题。不过相对应的，会比较麻烦，接下来介绍一下多机多卡DDP的使用方法。

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int) # 这个参数很重要
args = parser.parse_args()
torch.cuda.set_device(args.local_rank) # 调整计算的位置
# 以下二选一, 第一个是使用gloo后端需要设置的, 第二个是使用nccl需要设置的
os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'
# ps 检查nccl是否可用
# torch.distributed.is_nccl_available ()
torch.distributed.init_process_group(backend='nccl') # 选择nccl后端，初始化进程组
# 创建Dataloader
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, sampler=train_sampler)
# DDP进行训练
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

如何将PyTorch业务部署在K8S 集群中

• 构建docker 镜像

FROM pytorch/pytorch:2.2.1-cuda12.1-cudnn8-runtime

ADD . /opt/pytorch-mnist
WORKDIR /opt/pytorch-mnist

# Add folder for the logs.
RUN mkdir /katib

RUN chgrp -R 0 /opt/pytorch-mnist \
  && chmod -R g+rwX /opt/pytorch-mnist \
  && chgrp -R 0 /katib \
  && chmod -R g+rwX /katib

ENTRYPOINT ["python3", "/opt/pytorch-mnist/mnist.py"]

• 编写对应的yaml 文件（部署方式以kubeflow 为例）

apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: pytorch-results-pvc
spec:
  storageClassName: openebs-hostpath
  accessModes:
    - ReadWriteOnce
  resources:
    requests:
      storage: 10Gi
---
# 这里以方便处理的kubeflow 为例子
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple
  namespace: kubeflow
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch-mnist:2.2.1-cuda12.1-cudnn8-runtime
              imagePullPolicy: IfNotPresent
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=10"
                - "--batch-size"
                - "32"
                - "--test-batch-size"
                - "64"
                - "--lr"
                - "0.01"
                - "--momentum"
                - "0.9"
                - "--log-interval"
                - "10"
                - "--save-model"
                - "--log-path"
                - "/results/master.log"
              volumeMounts:
              - name: result-volume
                mountPath: /results
          volumes:
          - name: result-volume
            persistentVolumeClaim:
              claimName: pytorch-results-pvc
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: pytorch-mnist:2.2.1-cuda12.1-cudnn8-runtime
              imagePullPolicy: IfNotPresent
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=10"
                - "--batch-size"
                - "32"
                - "--test-batch-size"
                - "64"
                - "--lr"
                - "0.01"
                - "--momentum"
                - "0.9"
                - "--log-interval"
                - "10"
                - "--save-model"
                - "--log-path"
                - "/results/worker.log"
              volumeMounts:
              - name: result-volume
                mountPath: /results
          volumes:
          - name: result-volume
            persistentVolumeClaim:
              claimName: pytorch-results-pvc
      

后续

炼丹师的优雅内功三：kubeflow 如何承接机器学习任务

炼丹师的优雅内功四：volcano 批量处理容器计算业务

炼丹师的优雅内功五：大规模计算集群的秘密RDMA网络，IB网络和RoCe的纠葛