基于torch_dispatch机制生成Megatron-DeepSpeed调用关系图

2024-05-10 22:04

本文主要是介绍基于torch_dispatch机制生成Megatron-DeepSpeed调用关系图,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

基于torch_dispatch机制生成Megatron-DeepSpeed调用关系图

  • 一.局部效果图
  • 二.运行训练过程,拦截算子,生成调用关系信息
  • 三.可视化,生成SVG图像

想知道Megatron-DeepSpeed训练过程中各模块之间的调用关系。torch_dispatch机制可以拦截算子,inspect又能获取到调用栈(文件,类名,函数,行号).基于这些信息可以生成调用关系,最后用graphviz生成SVG图像。该思路也可以用来画其它pytorch工程的调用关系图

1.为了减少图像宽度,一行显示一级文件路径
2.没有显示具体的ATen算子。因为边太乱

一.局部效果图

在这里插入图片描述

二.运行训练过程,拦截算子,生成调用关系信息

# 前面构建模型的代码省略
from torch.utils._python_dispatch import TorchDispatchMode
import inspect
from dataclasses import dataclass
from typing import Any
import pickle@dataclass
class _ProfilerState:cls: Anyobject: Any = Noneclass TorchDumpDispatchMode(TorchDispatchMode):def __init__(self,parent):super().__init__()self.parent=parent        self.global_index=0        self.nodes=set()self.edges=set()def __del__(self):self.rank = torch.distributed.get_rank()graph={"nodes":self.nodes,"edges":self.edges}with open(f"call_graph_{self.rank}.pkl","wb") as f:pickle.dump(graph,f)def is_keep(self,node):# if node.function.find("wrapper")>=0:#     return False# if node.function.find("_call_impl")>=0:#     return Falsereturn Truedef __torch_dispatch__(self, func, types, args=(), kwargs=None):self.global_index+=1self.rank = torch.distributed.get_rank() func_packet = func._overloadpacket       if kwargs is None:kwargs = {}if self.rank==0:stacks=[i for i in inspect.stack() if self.is_keep(i)]stacks_sz=len(stacks)for idx in range(stacks_sz-1,1,-1):if "self" in stacks[idx].frame.f_locals:class_name = stacks[idx].frame.f_locals["self"].__class__.__name__else:class_name=""this_node=f"{stacks[idx].filename}:[{class_name}]:{stacks[idx].function}"if "self" in stacks[idx-1].frame.f_locals:class_name = stacks[idx-1].frame.f_locals["self"].__class__.__name__else:class_name=""                                    next_node=f"{stacks[idx-1].filename}:[{class_name}]:{stacks[idx-1].function}"self.nodes.add(this_node)self.nodes.add(next_node)self.edges.add(f"{this_node}->{next_node}")# if stacks_sz>1:#     if "self" in stacks[1].frame.f_locals:#         class_name = stacks[1].frame.f_locals["self"].__class__.__name__#     else:#         class_name=""                #     this_node=f"{stacks[1].filename}:[{class_name}]:{stacks[1].function}"#     next_node=f"{func_packet.__name__}"#     self.nodes.add(this_node)   #     self.nodes.add(next_node)            #     self.edges.add(f"{this_node}->{next_node}")ret= func(*args, **kwargs)return retclass TorchDumper:_CURRENT_Dumper = Nonedef __init__(self,schedule: Any):self.p= _ProfilerState(schedule)def __enter__(self):assert TorchDumper._CURRENT_Dumper is NoneTorchDumper._CURRENT_Dumper = selfif self.p.object is None:o = self.p.cls(self)o.__enter__()self.p.object = oelse:self.p.object.step()return selfdef __exit__(self, exc_type, exc_val, exc_tb):TorchDumper._CURRENT_Dumper = Noneif self.p.object is not None:self.p.object.__exit__(exc_type, exc_val, exc_tb)del self.p.object  #序列化保存def main():with TorchDumper(TorchDumpDispatchMode):#训练入口pretrain(train_valid_test_datasets_provider,model_provider,forward_step,extra_args_provider=llama_argument_handler,args_defaults={"tokenizer_type": "GPT2BPETokenizer"},)if __name__ == "__main__":main()

三.可视化,生成SVG图像

# coding=utf-8import os
from graphviz import Digraph,Graph
import pickle
import random
from distinctipy import distinctipydef generate_colors(N):'''生成N种有区别度的颜色'''result=[]for red, green, blue in distinctipy.get_colors(N):result.append("#{:02X}{:02X}{:02X}".format(int(red*255), int(green*255), int(blue*255)))return resultdef replace_name(name):'''修改节点名字(缩短,添加换行)'''if name.find("__torch_dispatch__")>=0:return Nonename=name.replace("/home/user/Megatron-DeepSpeed/","")name=name.replace("/home/anaconda3/envs/dev/lib/python3.10/site-packages/","")name=name.replace("/home/user/deepspeed/","")name=name.replace("/home/anaconda3/envs/dev/","")name=name.replace("/",r"\n")name=name.replace(":",r"\n")return name# 1.加载HOOK生成的调用关系文件
rank=0
with open(f"call_graph_{rank}.pkl","rb") as f:data=pickle.load(f)# 2.构建图,设置属性
dot = Digraph()
dot.node_attr = {"shape": "plaintext"}
dot.attr('graph', layout='dot')
dot.graph_attr.update(sep='4.0', ratio='compress')node_desc_id_map={}  #节点名与描述的关系映射表
src_node_color={}    #节点颜色映射表(同一个节点输出的边颜色一样)colors = generate_colors(10)
colors_sz=len(colors)fontsize="16"        #节点字体大小
penwidth="2.0"       #边宽度# 3.添加节点
for idx,v in enumerate(data["nodes"]):v=replace_name(v)if v is None:continuenode_desc_id_map[v]=f"{idx}"if v.find("megatron")>=0:dot.node(f"{idx}",v,style='filled',color='#73FBFD',fontsize=fontsize)elif v.find("deepspeed")>=0:dot.node(f"{idx}",v,style='filled',color='#FA8D89',fontsize=fontsize)else:dot.node(f"{idx}",v,style='filled',color='#C0C0C0',fontsize=fontsize)src_node_color[v]=colors[idx%colors_sz]# 4.添加边
for edge in data["edges"]:from_node,to_node=edge.split("->")from_node=replace_name(from_node)to_node=replace_name(to_node)if all([from_node,to_node]):color=src_node_color[from_node]dot.edge(node_desc_id_map[from_node], node_desc_id_map[to_node],color=color,penwidth=penwidth)# 5.保存SVG
save_path='megatron_deepspeed_callgraph'
dot.render(save_path,format='svg', view=False)# 6.修改背景色为灰色
import xml.etree.ElementTree as ET
svg_tree = ET.parse(f'{save_path}.svg')
root = svg_tree.getroot()
element = root.find(".//{http://www.w3.org/2000/svg}polygon")
element.set('fill', 'gray')
svg_tree.write(f'{save_path}.svg')

这篇关于基于torch_dispatch机制生成Megatron-DeepSpeed调用关系图的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/977679

相关文章

IDEA自动生成注释模板的配置教程

《IDEA自动生成注释模板的配置教程》本文介绍了如何在IntelliJIDEA中配置类和方法的注释模板,包括自动生成项目名称、包名、日期和时间等内容,以及如何定制参数和返回值的注释格式,需要的朋友可以... 目录项目场景配置方法类注释模板定义类开头的注释步骤类注释效果方法注释模板定义方法开头的注释步骤方法注

Python如何自动生成环境依赖包requirements

《Python如何自动生成环境依赖包requirements》:本文主要介绍Python如何自动生成环境依赖包requirements问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑... 目录生成当前 python 环境 安装的所有依赖包1、命令2、常见问题只生成当前 项目 的所有依赖包1、

SpringRetry重试机制之@Retryable注解与重试策略详解

《SpringRetry重试机制之@Retryable注解与重试策略详解》本文将详细介绍SpringRetry的重试机制,特别是@Retryable注解的使用及各种重试策略的配置,帮助开发者构建更加健... 目录引言一、SpringRetry基础知识二、启用SpringRetry三、@Retryable注解

如何在pycharm安装torch包

《如何在pycharm安装torch包》:本文主要介绍如何在pycharm安装torch包方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录在pycharm安装torch包适http://www.chinasem.cn配于我电脑的指令为适用的torch包为总结在p

SpringKafka错误处理(重试机制与死信队列)

《SpringKafka错误处理(重试机制与死信队列)》SpringKafka提供了全面的错误处理机制,通过灵活的重试策略和死信队列处理,下面就来介绍一下,具有一定的参考价值,感兴趣的可以了解一下... 目录引言一、Spring Kafka错误处理基础二、配置重试机制三、死信队列实现四、特定异常的处理策略五

MySQL中动态生成SQL语句去掉所有字段的空格的操作方法

《MySQL中动态生成SQL语句去掉所有字段的空格的操作方法》在数据库管理过程中,我们常常会遇到需要对表中字段进行清洗和整理的情况,本文将详细介绍如何在MySQL中动态生成SQL语句来去掉所有字段的空... 目录在mysql中动态生成SQL语句去掉所有字段的空格准备工作原理分析动态生成SQL语句在MySQL

pytorch之torch.flatten()和torch.nn.Flatten()的用法

《pytorch之torch.flatten()和torch.nn.Flatten()的用法》:本文主要介绍pytorch之torch.flatten()和torch.nn.Flatten()的用... 目录torch.flatten()和torch.nn.Flatten()的用法下面举例说明总结torch

Java调用C++动态库超详细步骤讲解(附源码)

《Java调用C++动态库超详细步骤讲解(附源码)》C语言因其高效和接近硬件的特性,时常会被用在性能要求较高或者需要直接操作硬件的场合,:本文主要介绍Java调用C++动态库的相关资料,文中通过代... 目录一、直接调用C++库第一步:动态库生成(vs2017+qt5.12.10)第二步:Java调用C++

Java利用docx4j+Freemarker生成word文档

《Java利用docx4j+Freemarker生成word文档》这篇文章主要为大家详细介绍了Java如何利用docx4j+Freemarker生成word文档,文中的示例代码讲解详细,感兴趣的小伙伴... 目录技术方案maven依赖创建模板文件实现代码技术方案Java 1.8 + docx4j + Fr

Java编译生成多个.class文件的原理和作用

《Java编译生成多个.class文件的原理和作用》作为一名经验丰富的开发者,在Java项目中执行编译后,可能会发现一个.java源文件有时会产生多个.class文件,从技术实现层面详细剖析这一现象... 目录一、内部类机制与.class文件生成成员内部类(常规内部类)局部内部类(方法内部类)匿名内部类二、