pytorch逻辑解析,第一层

这个过程对我个人而言其实有点过于复杂了,我到现在还没怎么理解,不过正是如此才需要写blog梳理一下思路和逻辑才能理解,我们从最基础最基础的东西开始一点一点向上捋:

首先,python是一门解释性语言,python脚本的每一句每一行语言都是要经过解释器从上到下逐句翻译的;其次所有的from和import都是从env环境中调取相应的函数来进行运算。

好,这两点是基础,我们在此基础之上逐句从上至下解释python脚本是怎么运行的、有哪些指令被launch了、指令是如何被执行的、指令又是如何调动数据的?数据在底层是如何被流转的?等等等等。

训练流程的三步

整个流程可以被分为三个步骤:

  1. 前向传播
  2. 反向传播
  3. 优化器更新

呃,对于有基础的同学这三步应该很熟悉了,前向传播就是在通过28个decoder层疯狂计算,算出一个巨型张量来表示每个输入值所对应下一个输出值的logits,然后再算出loss交叉熵;反向传播就是借助之前的中间激活值求梯度,偏导数计算,其实也就是按照计算图进行大量的矩阵乘法;优化器更新参数,这里用的是AdamW,稍稍复杂一点就是先更新一阶动量和二阶动量再去计算改变值。

上述三个步骤在我们的python脚本中其实很简单,三句代码:

outputs = model(
  input_ids=input_ids,
  attention_mask=attention_mask,
  labels=labels
)
loss = outputs.loss

loss.backward()

optimizer.step()

那么这三句代码是如何调动复杂的模型和数据进行更新的?

基本上可以分为四个阶段:动态架构解析与实例化、宏观控制流与构建状态图、ABI Boundary & Hardware Dispatch,四个阶段从顶层一点点下沉,每一个阶段每一个小步骤都需要执行也都需要花时间,每一小段时间都会在prefiling中体现,等我们在逻辑上捋清后再回头来分析具体的prefiling。

第一阶段:动态架构解析与实例化

第一阶段指的是:

model = AutoModelForCausalLM.from_pretrained(
  "./models/qwen/Qwen2.5-7B-Instruct",
  torch_dtype=torch.bfloat16
).to(device)

这段代码,AutoModelForCausalLM是一个工厂类,唯一的作用是接受一段字符串输入,并把Instruction Pointer指向字符串所代表的地址上去。

实际上这个Auto并没有在内存中创建任何一个对象,并不会让cpu执行任何实际的操作。PC跳转到文件读取函数,将/home/sysu/qwen_project/models/qwen/Qwen2___5-7B-Instruct/config.json加载到RAM中并且解析出architectures:

"architectures": [
  "Qwen2ForCausalLM"
]

transformer库本身就是hugging face工程师所维护的包含很多模型实体类代码放在一起以及所配套的一张维护好的注册表,PC执行到Qwen2ForCausalLM后就会跳转到Qwen2ForCausalLM实体类代码的物理地址然后开始执行,这个过程是Dynamic import动态导入和惰性加载lazy loading,也就是PC通过一层中介的注册表针对性的只将Qwen2ForCausalLM的实体类代码加载到RAM上,这个加载也就是系统的即时编译,JIT之后RAM中就有了Qwen2ForCausalLM的类代码,系统再在主存的划分出相应的物理空间来储存生成结构体,然后pc就可以读取到Qwen2ForCausalLM的真实物理地址,PC开始进入源码中依次去for循环执行如下的28层的前向传播。

好,现在我们的PC内现在储存了真正需要去执行的实体类代码,接下来我们需要把类变成实例,再通过指针相互嵌套链接而成层级结构的Python对象树。

这里需要补充一下,类是储存在RAM中的代码段,实例是系统根据类的设定在CPU中划分出空白物理内存块。现在回过头来,我们的PC上搭载的还是wen2ForCausalLM的默认执行的第一段机器码__init__,一句一句往下执行,重复着进行执行函数跳转、分配内存、写指针三个步骤:

class Qwen2Model:
  def __init__(self):
    self.num_layers = 28                  # 动作 A:写元数据
    self.layer_0 = Qwen2DecoderLayer()    # 动作 B & C:嵌套调用与记地址

class Qwen2DecoderLayer:
  def __init__(self):
    self.attention = Attention()          # 再次嵌套...

以这一段为例,CPU由于之前的__new__已经申请了一块空白内存,操作系统分配了一块内存和一个地址,此时root节点诞生,然后CPU带着这个地址开始跳入__init__代码段,先执行 self.num_layers = 28 这其实就是在向上面申请的空白内存存入数据,称之为实例的state;然后到下一句 self.layer_0 = Qwen2DecoderLayer() ,Qwen2DecoderLayer()是一个新的类名,此时CPU会暂存当前类的工作状态,然后再申请一块新的内存拿到新的地址,再去跳到新的地址上继续执行,就这样不断递归到最基层的算子如Linear、Embedding上,再往上不断回返,把调用栈里保存的进度弹出来,把子结点的地址存在父节点内部,循环往复,最后就把这个对象树构建出来了。

至今为止,这个对象树虽然层级很复杂,但它此时只包括了控制元数据和大量指针,我们下一步要做的就是把14GB大的7B模型搞到HBM2上并且将其挂载到这个对象树上。现代框架对此使用的是mmap机制,依然是AutoModelForCausalLM.from_pretrained这个超级庞大的函数,当对象树空壳搭建完成之后,PC会跳转到transformer库底层负责加载权重的模块,操作系统会在Python进程的虚拟地址空间中划出14GB的地址编号,这些地址在物理内存条上并不存在,因此主要的数据搬运工作由DMA控制器负责,DMA控制器接管PCIe总线跨过CPU把数据从硬盘直接copy到HBM2中。接下来要把空壳对象树和参数权重联系起来,用ATen,PyTorch的底层C++引擎,它会在CPU内存中生成一个极小的结构体专门储存14GB参数权重的起始物理地址,这个结构体会被包装成torch.nn.Parameter对象并被赋值给对象树的相应属性如self.weight属性上。

第一阶段,也是整个流程的起步阶段就此结束。