649 积分	0 好友	88 主题

发消息

PyTorch DeepSeek-V3计算图可视化实战：torchvista交互式探索指南

发表于 16 小时前 | 查看: 1| 回复: 0

在分析和理解大模型（LLM）时，我们常通过架构图来了解其网络结构（如下[1]）。然而，这类人工绘制的图示通常不会包含详尽的计算细节，例如张量（Tensor）的维度变化和具体操作过程。为深入分析数据从输入到输出的完整流程，本文将以DeepSeek-V3模型为例，介绍一种使用PyTorch进行网络计算图可视化的实用方法。

利用 torchvista [2] 工具，我们可以将 PyTorch 定义的 DeepSeek-V3 模型进行可视化，并在 Jupyter Notebook 等环境中获得一个可交互的计算流图，如下图所示。该可视化方法主要具备以下特点：

每个张量（Tensor）均会标注其尺寸信息；
每个模块可展开至底层的 PyTorch 基本操作；
采用类定义的模块支持折叠，展示深度可灵活调节；
可直接在界面中查看每个操作的具体定义。

示例地址（若无法访问，请参考文末的 GitHub 地址）：
https://colab.research.google.com/drive/1q0NAAnCP_e0c53wacM5H0cqDzguqTOER?usp=sharing

下面，我们将分步介绍如何构建此交互式计算图。

一、torchvista 基本使用

torchvista 是一款交互式可视化工具，可直接在 Notebook 环境中对 PyTorch 模型（nn.Module）的前向传播过程进行可视化。它非常适用于辅助理解模型推理时的详细计算流程。

首先进行工具安装：

!pip install --upgrade torchvista

我们构建一个简单的线性层模型来测试其功能：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvista import trace_model

# 模型定义：
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
    def forward(self, x):
        return self.linear(x)

# 创建模型、定义输入
model = LinearModel()
example_input = torch.randn(2, 10)

# 调用 torchvista 来可视化模型
trace_model(
    model,  # 模型
    example_input, # 输入
    collapse_modules_after_depth=3,
    show_non_gradient_nodes=False,
    forced_module_tracing_depth=None,
    height=500,
    export_format=None
)

输出结果如下图所示：

trace_model 函数的关键参数说明：

collapse_modules_after_depth：模块自动折叠的深度，默认为0。
show_non_gradient_nodes：是否显示不参与梯度计算的节点。
forced_module_tracing_depth：控制追踪（trace）的层数。设置较小的值可以降低图形渲染的计算成本，默认值为None。
height：Notebook中图形显示框的高度。
export_format：支持导出为 'png'、'svg'、'html' 格式。

实用提示：

当生成交互图速度较慢时，可尝试降低 forced_module_tracing_depth 参数的值。
如果遇到非 nn.Module 模块显示错误，可尝试将其替换为对应的 PyTorch 底层操作。

二、DeepSeek-V3模型的可视化实战

使用 PyTorch 构建 DeepSeek-V3 模型会涉及到量化、FP8计算等技术，可参考官方文档[3]进行模型搭建。由于模型依赖一些 GPU kernel，首先需要在 Notebook 中安装 triton 库：

# 安装 triton
!pip install --upgrade triton

然后，参照官方代码（参考：https://huggingface.co/deepseek-ai/DeepSeek-V3）搭建对应的模型。我们首先尝试可视化一个包含 MLA 和 FFN 层的 Block 的计算过程。

args = ModelArgs()
model = Block(0, args)
torch.set_default_dtype(torch.bfloat16)
tokens = torch.randn((1, 1024, 2048) , dtype=torch.bfloat16)
freqs_cis = precompute_freqs_cis(args)
seqlen = tokens.size(1)
freqs_cis = freqs_cis[0:seqlen]
mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device).triu_(1)
trace_model(model, (tokens,0, freqs_cis, mask), export_format='png')

Block 中 MLA 子模块的计算流图如下：

为了观察模型整体结构，我们将模型总层数修改为2：

n_layers: int = 2   # 总层数 -> 2
n_dense_layers: int = 1

然后运行以下代码进行可视化：

args = ModelArgs()
model = Transformer(args)
torch.set_default_dtype(torch.bfloat16)
x = torch.randint(0, args.vocab_size, (2, 128)) # batch_size: 2, sequence_length:128
trace_model(model, (x,))

获得的计算流图如下图所示，可以清晰地看到两个Block中FFN层的主要差异：一个使用标准的MLP，另一个则采用了MoE（混合专家）结构。