尧图网站建设 
文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载导读本文基于 pysheeet 仓库 src/megatron 目录下的完整工具链讲解如何用 Megatron Bridge 的 recipe 式接口在 SLURM EFAElastic Fabric Adapter集群上启动 Megatron-LM 预训练任务从 Docker 镜像构建、enroot 镜像导出、srun.sh启动、Hydra 风格参数覆盖到 Nsys 与 VizTracer 两套剖析方案的启用方法。读完本文你将掌握一套可直接复制的镜像构建 → 分布式启动 → 性能剖析完整流程并理解 Megatron Bridge 为何选择python直启而非torchrun以及如何通过 monkey-patch 扩展自定义 profiler。关联文档docs/notes/llm/megatron.rst相关进阶话题见 docs/notes/appendix/megatron-efa-monitoring.rst。背景Megatron-LM 与 Megatron Bridge 的分工Megatron-LM 是 NVIDIA 用于训练和微调大型 Transformer 模型的框架原生支持张量并行tensor parallelism、流水线并行pipeline parallelism与专家并行expert parallelism。Megatron Bridge 则位于 Megatron-LM 之上提供基于 recipe 的配置接口——不再需要手工拼接几十个 CLI 标志只需编写一个返回配置对象的短小 Python recipe 文件即可。Recipe 最大的价值在于可以通过hf_path参数直接加载 HuggingFace 预训练权重无需手动做 checkpoint 格式转换即可从任意 HuggingFace 模型起步开始训练。这一点正是 pysheeet 仓库 src/megatron/recipes 下 DeepSeek V2 Lite 预训练示例的核心用法。环境准备从 Docker 构建到 enroot 镜像导出Megatron 训练环境被打包为 Docker 镜像再经 enroot 导出为 squashfs.sqsh文件。enroot 是面向 HPC 的轻量容器运行时它能把 Docker 镜像转换成无特权沙箱并借助 pyxis 插件与 SLURM 集成。当srun.sh运行时会把--container-image与--container-mounts传给srun由 pyxis 负责导入.sqsh并在容器内启动每个任务。构建镜像在 src/megatron 目录下执行make build该命令通过 Makefile 调用 enroot.shIMAGE_NAME ? megatron-lm build: ./enroot.sh -n $(IMAGE_NAME) -f Dockerfileenroot.sh依次完成两步先用docker build -f Dockerfile -t megatron-lm .构建镜像再用enroot import -o megatron-lmlatest.sqsh dockerd://megatron-lm导出为 enroot 镜像。构建产物megatron-lmlatest.sqsh会出现在当前目录之后 srun.sh 会通过SQSH环境变量默认./megatron-lmlatest.sqsh自动拾取它。镜像内部组件Dockerfile 展示了镜像的完整技术栈这些组件与后续 profiling 与 EFA 网络调优直接相关组件版本参数用途CUDACUDA_VERSION12.8.1基础镜像devel-ubuntu24.04GDRCopyGDRCOPY_VERSIONv2.5.1GPU 直接内存复制EFA 安装器EFA_INSTALLER_VERSION1.47.0安装 libfabric/EFA 驱动aws-ofi-nccl固定 commitNCCL 的 AWS OFI 网络插件含 GIN 支持NCCLNCCL_VERSIONv2.29.3-1从源码编译含 Device API 支持NCCL-testsNCCL_TESTS_VERSIONv2.17.9网络验证工具NVSHMEMNVSHMEM_VERSIONv3.5.19-1对称内存通信库PyTorchTORCH_VERSION2.9.1训练框架Megatron-BridgeMEGATRON_BRIDGE_VERSIONv0.2.2recipe 配置接口镜像内还安装了 Nsight Systems CLInsight-systems-cli与viztracer这是后文 Nsys 与 VizTracer 剖析的前提。构建完成后 Dockerfile 末尾有 sanity check验证torch、transformer_engine、megatron.core与megatron.bridge均能正常导入。覆盖镜像路径与挂载make build之外可通过环境变量覆盖启动时的镜像路径与容器挂载变量默认值说明SQSH./megatron-lmlatest.sqshenroot 镜像路径MOUNT.:/workspace/megatron,/fsx:/fsx容器挂载项GPUS_PER_NODE8每节点 GPU 数其中MOUNT的默认值把仓库目录挂载到容器内/workspace/megatron并把宿主机/fsx预训练权重通常放在这里挂载到容器内同名路径。Recipe用 Python 配置对象替代命令行参数Recipe 是一个 Python 文件调用 Megatron Bridge 的配置函数并返回ConfigContainer。仓库示例 src/megatron/recipes/deepseek_v2_lite_pretrain.py 完整如下from megatron.bridge.recipes.deepseek.deepseek_v2 import ( deepseek_v2_lite_pretrain_config, ) def configure(hf_pathNone, moe_token_dispatcher_typeNone): cfg deepseek_v2_lite_pretrain_config( **({hf_path: hf_path} if hf_path else {}), tensor_model_parallel_size8, pipeline_model_parallel_size1, expert_model_parallel_size2, sequence_parallelTrue, seq_length4096, train_iters500, global_batch_size64, micro_batch_size1, eval_interval100, lr_warmup_iters50, save_interval0, ) cfg.model.moe_permute_fusion False if moe_token_dispatcher_type deepep: cfg.model.moe_token_dispatcher_type flex cfg.model.moe_flex_dispatcher_backend deepep cfg.model.moe_enable_deepep True cfg.model.moe_shared_expert_overlap False return cfg关键点解读并行配置tensor_model_parallel_size8、pipeline_model_parallel_size1、expert_model_parallel_size2组合在一起8 卡节点上每张卡承担一个 TP 分片同时启用sequence_parallelTrue降低激活显存占用。hf_path参数为 None 时不传入即从随机初始化训练传入时直接从 HuggingFace 路径加载预训练权重这正是免手动转换的核心能力。训练规模seq_length4096、global_batch_size64、micro_batch_size1流水线并行场景下 microbatch 较小属正常设计、train_iters500、lr_warmup_iters50。MoE 开关cfg.model.moe_permute_fusion False关闭 permute 融合当moe_token_dispatcher_type deepep时切换为 flex dispatcher 后端 DeepEP并关闭共享专家重叠。这是在 recipe 层面对 DeepSeek MoE 模型启用 DeepEP含 NCCL GIN的全部配置。启动预训练srun.sh 的完整工作流2 节点 DeepSeek V2 Lite 预训练salloc -N 2 ./srun.sh recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-LiteHydra 风格参数覆盖不需要改 recipe 文件直接在命令行追加keyvalue即可覆盖任意配置项./srun.sh recipes/deepseek_v2_lite_pretrain.py \ train.train_iters1000train.train_iters1000会把训练迭代数从 recipe 里的 500 覆盖为 1000。这一机制由 entrypoint.py 实现它先按参数解析规则抽取hf_path与moe_token_dispatcher_type两个特殊参数传给configure()其余参数走parse_cli_overrides()——通过create_omegaconf_dict_config把ConfigContainer转成 OmegaConf 字典再经parse_hydra_overrides解析keyvalue最后由apply_overrides写回配置对象最后调用pretrain(configcfg, forward_step_funcforward_step)启动训练。srun.sh 内部做了什么srun.sh 是整个启动链路的核心其职责可分为三层解析参数识别--nsys标志与 recipe 路径、剩余 overrides。组装容器内命令通过scontrol show hostnames与getent hosts解析MASTER_ADDR随后导出 EFA/NCCL 环境变量并执行srun启动。环境变量映射把 SLURM 提供的SLURM_PROCID/SLURM_LOCALID/SLURM_NTASKS分别映射为 Megatron 分布式初始化所需的RANK/LOCAL_RANK/WORLD_SIZE这样无需torchrun即可完成分布式初始化。srun.sh 内设置的关键环境变量如下与 Dockerfile 中的镜像级默认值相互印证export LD_LIBRARY_PATH/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:$LD_LIBRARY_PATH export FI_PROVIDERefa export FI_EFA_USE_DEVICE_RDMA1 export FI_EFA_FORK_SAFE1 export NCCL_NET_PLUGIN/opt/amazon/ofi-nccl/lib/libnccl-net-ofi.so export NCCL_TUNER_PLUGIN/opt/amazon/ofi-nccl/lib/libnccl-tuner-ofi.so export NCCL_DEBUGWARN export NCCL_BUFFSIZE8388608 export NCCL_P2P_NET_CHUNKSIZE524288 export OMP_NUM_THREADS1 export TRITON_CACHE_DIR/tmp/triton_cache_${SLURM_PROCID:-0} export TORCH_EXTENSIONS_DIR/tmp/torch_ext_${SLURM_PROCID:-0} export MASTER_ADDR${master_addr} export MASTER_PORT29500 export DEEP_EP_BACKENDnccl export NCCL_GIN_TYPE2 export CUDA_DEVICE_MAX_CONNECTIONS1 export TORCH_DISTRIBUTED_BACKENDnccl export RANK${SLURM_PROCID} export LOCAL_RANK${SLURM_LOCALID} export WORLD_SIZE${SLURM_NTASKS}最终通过 pyxis 插件启动srun --container-image ${SQSH} \ --container-mounts ${MOUNT} \ --container-name megatron \ --mpipmix \ --ntasks-per-node${GPUS_PER_NODE} \ bash -c ${cmd}注意--mpipmix与--ntasks-per-node${GPUS_PER_NODE}每个 GPU 对应一个进程由 SLURM 直接管理而不是由torchrun派生。启用 Nsys 剖析精确捕获指定 step 区间命令行用法./srun.sh --nsys recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-Lite \ profiling.use_nsys_profilertrue \ profiling.profile_step_start10 \ profiling.profile_step_end15 \ profiling.profile_ranks[0]底层行为当--nsys开启且当前进程RANK0时srun.sh 会拼装如下 nsys 命令并前置到python3 entrypoint.py之前nsys profile \ -t cuda,nvtx \ -s none \ --cpuctxswnone \ --capture-rangecudaProfilerApi \ --capture-range-endstop \ --enable efa_metrics \ -o /workspace/megatron/nsys-megatron/profile-$(hostname)-rank${RANK}.nsys-rep \ --force-overwritetrue其中--capture-rangecudaProfilerApi以cudaProfilerStart/cudaProfilerStopAPI 作为捕获边界因此只有profile_step_start到profile_step_end之间的训练 step 会被捕获默认 1015避免剖析整个训练过程产生过大的 trace。--enable efa_metrics在 AWS 上监控 EFA 网络流量的关键开关。Nsys 会按约 10Hz 采样 EFA 设备计数器如rdmap113s0、rdmap114s0采集 TX/RX 带宽、收发包数与错误计数并利用 aws-ofi-nccl 的 NVTX 注解把 NCCL 集合通信与 EFA 网络活动关联到同一条时间线从而定位通信瓶颈。输出路径.nsys-rep文件写入容器挂载内的nsys-megatron/目录即宿主机上的src/megatron/nsys-megatron/可用 Nsight Systems GUI 打开分析。关于 EFA 网络监控的完整实战含 NCCL GIN 启用与日志验证请参考仓库内配套文章 docs/notes/appendix/megatron-efa-monitoring.rst。为什么用python直启而不是torchrunsrun.sh 直接启动python3 entrypoint.py刻意不用torchrun这是有意的设计决策torchrun通过multiprocessing派生子进程spawn 边界会干扰 Nsys 剖析——profiler 有时无法捕获子进程发出的cudaProfilerStart/cudaProfilerStop调用导致 trace 为空或不完整。在srun --ntasks-per-nodeGPUS模式下每张 GPU 拥有一个由 SLURM 直接管理的进程RANK、LOCAL_RANK、WORLD_SIZE分别由SLURM_PROCID、SLURM_LOCALID、SLURM_NTASKS推导而来完全绕开 spawn 层让 Nsys以及 VizTracer 等其他 profiler获得每个 rank 干净、单进程的视图。自定义 Profiler以 VizTracer 插件为例Megatron Bridge 的 profiling hooks 是可扩展的。src/megatron/viztracer_plugin.py 展示了如何通过 monkey-patchmegatron.bridge.training.profiling接入自定义 profiler注册新配置字段用dataclasses.field把use_viztracer加入ProfilingConfig的__dataclass_fields__使 OmegaConf 认可该 override。patchhandle_profiling_step在profile_step_start迭代处启动 VizTracer先委托原函数处理 nsys/pytorch profiler若未启用再检查use_viztracer、should_profile_rank与 step 匹配。patchhandle_profiling_stop在profile_step_end迭代处调用stop()与save()把 JSON trace 写入viztracer_output/trace_rank{rank}.json输出目录可用环境变量VIZTRACER_OUTPUT_DIR覆盖。互斥校验patch 后的finalize()断言use_viztracer与use_nsys_profiler、use_pytorch_profiler互斥避免同时启用多个 profiler。插件在 entrypoint.py 中于任何 Megatron import 之前加载import viztracer_plugin viztracer_plugin.install()命令行启用 VizTracer./srun.sh recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-Lite \ profiling.use_viztracertrue \ profiling.profile_step_start10 \ profiling.profile_step_end15 \ profiling.profile_ranks[0]生成的 JSON trace 可用 VizTracer Web UI 或 Chrome 的chrome://tracing打开。由于log_torchTrue还能看到 NCCL stream 与 CUDA stream 等 PyTorch 层细节。注意若要观察 EFA 适配器级别的带宽/包数/错误计数仍需使用 Nsys 的--enable efa_metrics。这套模式对任何 profiler 通用——只需在 patch 的 hooks 里实现start()/stop()/save()并注册一个配置字段即可通过 Hydra override 开关控制。常见问题与排查要点镜像未找到确认make build已生成megatron-lmlatest.sqsh或通过SQSH环境变量显式指定路径。权重路径hf_path/fsx/models/...依赖MOUNT中的/fsx:/fsx挂载需确保宿主机上权重文件真实存在。覆盖无效确保 override 的 key 与 OmegaConf 结构一致如train.train_iters1000、profiling.profile_step_start10新增 profiler 字段需先注册进ProfilingConfig否则 OmegaConf 会拒绝该 override。GIN 生效验证DeepEP NCCL GIN 模式下训练日志应出现[NCCL] Device API initialized、[NCCL] GIN proxy mode enabled (type2)等初始化信息详见 docs/notes/appendix/megatron-efa-monitoring.rst。小结pysheeet 的 src/megatron 目录提供了一条完整的 Megatron-LM 预训练实战链路make build构建 Docker 并导出 enroot 镜像 → recipe 声明式配置支持hf_path直接加载 HuggingFace 权重与 DeepEP MoE 调度→srun.sh pyxis 在 SLURM/EFA 上以每 GPU 一个进程方式启动 →--nsys或use_viztracer精确剖析指定训练 step。其中python直启替代torchrun、monkey-patch profiling hooks 接入自定义 profiler 的设计为大规模分布式训练的性能剖析提供了可复用的工程范式。赞分享文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载相关推荐使用 NCCL GIN、Nsys EFA 指标与 Viztracer 监控 Megatron-LM 分布式训练网络性能使用 NCCL GIN、Nsys EFA 指标与 Viztracer 监控 Megatron LM 分布式训练网络性能 在 AWS 上使用 Megatron L文档教程开发工具CogVideoX高性能训练实战Slurm与PBS集群调度终极指南CogVideoX高性能训练实战Slurm与PBS集群调度终极指南 CogVideoX作为领先的文本与图像到视频生成模型其高性能训练需要强大的集群调度支持。人工智能大模型媒体生成预训练微调如何在 Slurm 集群上运行 Faiss 分布式 k-means 训练大量质心如何在 Slurm 集群上运行 Faiss 分布式 k means 训练大量质心 当你需要给上亿向量级的训练集聚出百万级质心例如用 IVF 索引切 100机器学习搜索引擎向量数据库上一篇5分钟掌握WeTextProcessing一站式文本标准化处理的开源利器下一篇OpenCVE部署架构详解Docker容器化部署和配置最佳实践创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 使用 Megatron Bridge 在 SLURM + EFA 集群上运行 DeepSeek 预训练与 Nsys 性能剖析（pysheeet 实战指南）
文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载导读本文基于 pysheeet 仓库 src/megatron 目录下的完整工具链讲解如何用 Megatron Bridge 的 recipe 式接口在 SLURM EFAElastic Fabric Adapter集群上启动 Megatron-LM 预训练任务从 Docker 镜像构建、enroot 镜像导出、srun.sh启动、Hydra 风格参数覆盖到 Nsys 与 VizTracer 两套剖析方案的启用方法。读完本文你将掌握一套可直接复制的镜像构建 → 分布式启动 → 性能剖析完整流程并理解 Megatron Bridge 为何选择python直启而非torchrun以及如何通过 monkey-patch 扩展自定义 profiler。关联文档docs/notes/llm/megatron.rst相关进阶话题见 docs/notes/appendix/megatron-efa-monitoring.rst。背景Megatron-LM 与 Megatron Bridge 的分工Megatron-LM 是 NVIDIA 用于训练和微调大型 Transformer 模型的框架原生支持张量并行tensor parallelism、流水线并行pipeline parallelism与专家并行expert parallelism。Megatron Bridge 则位于 Megatron-LM 之上提供基于 recipe 的配置接口——不再需要手工拼接几十个 CLI 标志只需编写一个返回配置对象的短小 Python recipe 文件即可。Recipe 最大的价值在于可以通过hf_path参数直接加载 HuggingFace 预训练权重无需手动做 checkpoint 格式转换即可从任意 HuggingFace 模型起步开始训练。这一点正是 pysheeet 仓库 src/megatron/recipes 下 DeepSeek V2 Lite 预训练示例的核心用法。环境准备从 Docker 构建到 enroot 镜像导出Megatron 训练环境被打包为 Docker 镜像再经 enroot 导出为 squashfs.sqsh文件。enroot 是面向 HPC 的轻量容器运行时它能把 Docker 镜像转换成无特权沙箱并借助 pyxis 插件与 SLURM 集成。当srun.sh运行时会把--container-image与--container-mounts传给srun由 pyxis 负责导入.sqsh并在容器内启动每个任务。构建镜像在 src/megatron 目录下执行make build该命令通过 Makefile 调用 enroot.shIMAGE_NAME ? megatron-lm build: ./enroot.sh -n $(IMAGE_NAME) -f Dockerfileenroot.sh依次完成两步先用docker build -f Dockerfile -t megatron-lm .构建镜像再用enroot import -o megatron-lmlatest.sqsh dockerd://megatron-lm导出为 enroot 镜像。构建产物megatron-lmlatest.sqsh会出现在当前目录之后 srun.sh 会通过SQSH环境变量默认./megatron-lmlatest.sqsh自动拾取它。镜像内部组件Dockerfile 展示了镜像的完整技术栈这些组件与后续 profiling 与 EFA 网络调优直接相关组件版本参数用途CUDACUDA_VERSION12.8.1基础镜像devel-ubuntu24.04GDRCopyGDRCOPY_VERSIONv2.5.1GPU 直接内存复制EFA 安装器EFA_INSTALLER_VERSION1.47.0安装 libfabric/EFA 驱动aws-ofi-nccl固定 commitNCCL 的 AWS OFI 网络插件含 GIN 支持NCCLNCCL_VERSIONv2.29.3-1从源码编译含 Device API 支持NCCL-testsNCCL_TESTS_VERSIONv2.17.9网络验证工具NVSHMEMNVSHMEM_VERSIONv3.5.19-1对称内存通信库PyTorchTORCH_VERSION2.9.1训练框架Megatron-BridgeMEGATRON_BRIDGE_VERSIONv0.2.2recipe 配置接口镜像内还安装了 Nsight Systems CLInsight-systems-cli与viztracer这是后文 Nsys 与 VizTracer 剖析的前提。构建完成后 Dockerfile 末尾有 sanity check验证torch、transformer_engine、megatron.core与megatron.bridge均能正常导入。覆盖镜像路径与挂载make build之外可通过环境变量覆盖启动时的镜像路径与容器挂载变量默认值说明SQSH./megatron-lmlatest.sqshenroot 镜像路径MOUNT.:/workspace/megatron,/fsx:/fsx容器挂载项GPUS_PER_NODE8每节点 GPU 数其中MOUNT的默认值把仓库目录挂载到容器内/workspace/megatron并把宿主机/fsx预训练权重通常放在这里挂载到容器内同名路径。Recipe用 Python 配置对象替代命令行参数Recipe 是一个 Python 文件调用 Megatron Bridge 的配置函数并返回ConfigContainer。仓库示例 src/megatron/recipes/deepseek_v2_lite_pretrain.py 完整如下from megatron.bridge.recipes.deepseek.deepseek_v2 import ( deepseek_v2_lite_pretrain_config, ) def configure(hf_pathNone, moe_token_dispatcher_typeNone): cfg deepseek_v2_lite_pretrain_config( **({hf_path: hf_path} if hf_path else {}), tensor_model_parallel_size8, pipeline_model_parallel_size1, expert_model_parallel_size2, sequence_parallelTrue, seq_length4096, train_iters500, global_batch_size64, micro_batch_size1, eval_interval100, lr_warmup_iters50, save_interval0, ) cfg.model.moe_permute_fusion False if moe_token_dispatcher_type deepep: cfg.model.moe_token_dispatcher_type flex cfg.model.moe_flex_dispatcher_backend deepep cfg.model.moe_enable_deepep True cfg.model.moe_shared_expert_overlap False return cfg关键点解读并行配置tensor_model_parallel_size8、pipeline_model_parallel_size1、expert_model_parallel_size2组合在一起8 卡节点上每张卡承担一个 TP 分片同时启用sequence_parallelTrue降低激活显存占用。hf_path参数为 None 时不传入即从随机初始化训练传入时直接从 HuggingFace 路径加载预训练权重这正是免手动转换的核心能力。训练规模seq_length4096、global_batch_size64、micro_batch_size1流水线并行场景下 microbatch 较小属正常设计、train_iters500、lr_warmup_iters50。MoE 开关cfg.model.moe_permute_fusion False关闭 permute 融合当moe_token_dispatcher_type deepep时切换为 flex dispatcher 后端 DeepEP并关闭共享专家重叠。这是在 recipe 层面对 DeepSeek MoE 模型启用 DeepEP含 NCCL GIN的全部配置。启动预训练srun.sh 的完整工作流2 节点 DeepSeek V2 Lite 预训练salloc -N 2 ./srun.sh recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-LiteHydra 风格参数覆盖不需要改 recipe 文件直接在命令行追加keyvalue即可覆盖任意配置项./srun.sh recipes/deepseek_v2_lite_pretrain.py \ train.train_iters1000train.train_iters1000会把训练迭代数从 recipe 里的 500 覆盖为 1000。这一机制由 entrypoint.py 实现它先按参数解析规则抽取hf_path与moe_token_dispatcher_type两个特殊参数传给configure()其余参数走parse_cli_overrides()——通过create_omegaconf_dict_config把ConfigContainer转成 OmegaConf 字典再经parse_hydra_overrides解析keyvalue最后由apply_overrides写回配置对象最后调用pretrain(configcfg, forward_step_funcforward_step)启动训练。srun.sh 内部做了什么srun.sh 是整个启动链路的核心其职责可分为三层解析参数识别--nsys标志与 recipe 路径、剩余 overrides。组装容器内命令通过scontrol show hostnames与getent hosts解析MASTER_ADDR随后导出 EFA/NCCL 环境变量并执行srun启动。环境变量映射把 SLURM 提供的SLURM_PROCID/SLURM_LOCALID/SLURM_NTASKS分别映射为 Megatron 分布式初始化所需的RANK/LOCAL_RANK/WORLD_SIZE这样无需torchrun即可完成分布式初始化。srun.sh 内设置的关键环境变量如下与 Dockerfile 中的镜像级默认值相互印证export LD_LIBRARY_PATH/opt/amazon/ofi-nccl/lib:/opt/amazon/efa/lib:/opt/aws-ofi-nccl/install/lib:$LD_LIBRARY_PATH export FI_PROVIDERefa export FI_EFA_USE_DEVICE_RDMA1 export FI_EFA_FORK_SAFE1 export NCCL_NET_PLUGIN/opt/amazon/ofi-nccl/lib/libnccl-net-ofi.so export NCCL_TUNER_PLUGIN/opt/amazon/ofi-nccl/lib/libnccl-tuner-ofi.so export NCCL_DEBUGWARN export NCCL_BUFFSIZE8388608 export NCCL_P2P_NET_CHUNKSIZE524288 export OMP_NUM_THREADS1 export TRITON_CACHE_DIR/tmp/triton_cache_${SLURM_PROCID:-0} export TORCH_EXTENSIONS_DIR/tmp/torch_ext_${SLURM_PROCID:-0} export MASTER_ADDR${master_addr} export MASTER_PORT29500 export DEEP_EP_BACKENDnccl export NCCL_GIN_TYPE2 export CUDA_DEVICE_MAX_CONNECTIONS1 export TORCH_DISTRIBUTED_BACKENDnccl export RANK${SLURM_PROCID} export LOCAL_RANK${SLURM_LOCALID} export WORLD_SIZE${SLURM_NTASKS}最终通过 pyxis 插件启动srun --container-image ${SQSH} \ --container-mounts ${MOUNT} \ --container-name megatron \ --mpipmix \ --ntasks-per-node${GPUS_PER_NODE} \ bash -c ${cmd}注意--mpipmix与--ntasks-per-node${GPUS_PER_NODE}每个 GPU 对应一个进程由 SLURM 直接管理而不是由torchrun派生。启用 Nsys 剖析精确捕获指定 step 区间命令行用法./srun.sh --nsys recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-Lite \ profiling.use_nsys_profilertrue \ profiling.profile_step_start10 \ profiling.profile_step_end15 \ profiling.profile_ranks[0]底层行为当--nsys开启且当前进程RANK0时srun.sh 会拼装如下 nsys 命令并前置到python3 entrypoint.py之前nsys profile \ -t cuda,nvtx \ -s none \ --cpuctxswnone \ --capture-rangecudaProfilerApi \ --capture-range-endstop \ --enable efa_metrics \ -o /workspace/megatron/nsys-megatron/profile-$(hostname)-rank${RANK}.nsys-rep \ --force-overwritetrue其中--capture-rangecudaProfilerApi以cudaProfilerStart/cudaProfilerStopAPI 作为捕获边界因此只有profile_step_start到profile_step_end之间的训练 step 会被捕获默认 1015避免剖析整个训练过程产生过大的 trace。--enable efa_metrics在 AWS 上监控 EFA 网络流量的关键开关。Nsys 会按约 10Hz 采样 EFA 设备计数器如rdmap113s0、rdmap114s0采集 TX/RX 带宽、收发包数与错误计数并利用 aws-ofi-nccl 的 NVTX 注解把 NCCL 集合通信与 EFA 网络活动关联到同一条时间线从而定位通信瓶颈。输出路径.nsys-rep文件写入容器挂载内的nsys-megatron/目录即宿主机上的src/megatron/nsys-megatron/可用 Nsight Systems GUI 打开分析。关于 EFA 网络监控的完整实战含 NCCL GIN 启用与日志验证请参考仓库内配套文章 docs/notes/appendix/megatron-efa-monitoring.rst。为什么用python直启而不是torchrunsrun.sh 直接启动python3 entrypoint.py刻意不用torchrun这是有意的设计决策torchrun通过multiprocessing派生子进程spawn 边界会干扰 Nsys 剖析——profiler 有时无法捕获子进程发出的cudaProfilerStart/cudaProfilerStop调用导致 trace 为空或不完整。在srun --ntasks-per-nodeGPUS模式下每张 GPU 拥有一个由 SLURM 直接管理的进程RANK、LOCAL_RANK、WORLD_SIZE分别由SLURM_PROCID、SLURM_LOCALID、SLURM_NTASKS推导而来完全绕开 spawn 层让 Nsys以及 VizTracer 等其他 profiler获得每个 rank 干净、单进程的视图。自定义 Profiler以 VizTracer 插件为例Megatron Bridge 的 profiling hooks 是可扩展的。src/megatron/viztracer_plugin.py 展示了如何通过 monkey-patchmegatron.bridge.training.profiling接入自定义 profiler注册新配置字段用dataclasses.field把use_viztracer加入ProfilingConfig的__dataclass_fields__使 OmegaConf 认可该 override。patchhandle_profiling_step在profile_step_start迭代处启动 VizTracer先委托原函数处理 nsys/pytorch profiler若未启用再检查use_viztracer、should_profile_rank与 step 匹配。patchhandle_profiling_stop在profile_step_end迭代处调用stop()与save()把 JSON trace 写入viztracer_output/trace_rank{rank}.json输出目录可用环境变量VIZTRACER_OUTPUT_DIR覆盖。互斥校验patch 后的finalize()断言use_viztracer与use_nsys_profiler、use_pytorch_profiler互斥避免同时启用多个 profiler。插件在 entrypoint.py 中于任何 Megatron import 之前加载import viztracer_plugin viztracer_plugin.install()命令行启用 VizTracer./srun.sh recipes/deepseek_v2_lite_pretrain.py \ hf_path/fsx/models/deepseek-ai/DeepSeek-V2-Lite \ profiling.use_viztracertrue \ profiling.profile_step_start10 \ profiling.profile_step_end15 \ profiling.profile_ranks[0]生成的 JSON trace 可用 VizTracer Web UI 或 Chrome 的chrome://tracing打开。由于log_torchTrue还能看到 NCCL stream 与 CUDA stream 等 PyTorch 层细节。注意若要观察 EFA 适配器级别的带宽/包数/错误计数仍需使用 Nsys 的--enable efa_metrics。这套模式对任何 profiler 通用——只需在 patch 的 hooks 里实现start()/stop()/save()并注册一个配置字段即可通过 Hydra override 开关控制。常见问题与排查要点镜像未找到确认make build已生成megatron-lmlatest.sqsh或通过SQSH环境变量显式指定路径。权重路径hf_path/fsx/models/...依赖MOUNT中的/fsx:/fsx挂载需确保宿主机上权重文件真实存在。覆盖无效确保 override 的 key 与 OmegaConf 结构一致如train.train_iters1000、profiling.profile_step_start10新增 profiler 字段需先注册进ProfilingConfig否则 OmegaConf 会拒绝该 override。GIN 生效验证DeepEP NCCL GIN 模式下训练日志应出现[NCCL] Device API initialized、[NCCL] GIN proxy mode enabled (type2)等初始化信息详见 docs/notes/appendix/megatron-efa-monitoring.rst。小结pysheeet 的 src/megatron 目录提供了一条完整的 Megatron-LM 预训练实战链路make build构建 Docker 并导出 enroot 镜像 → recipe 声明式配置支持hf_path直接加载 HuggingFace 权重与 DeepEP MoE 调度→srun.sh pyxis 在 SLURM/EFA 上以每 GPU 一个进程方式启动 →--nsys或use_viztracer精确剖析指定训练 step。其中python直启替代torchrun、monkey-patch profiling hooks 接入自定义 profiler 的设计为大规模分布式训练的性能剖析提供了可复用的工程范式。赞分享文档教程开发工具【免费下载链接】pysheeetPython Cheat Sheet项目地址https://gitcode.com/gh_mirrors/py/pysheeet点击查看免费下载相关推荐使用 NCCL GIN、Nsys EFA 指标与 Viztracer 监控 Megatron-LM 分布式训练网络性能使用 NCCL GIN、Nsys EFA 指标与 Viztracer 监控 Megatron LM 分布式训练网络性能 在 AWS 上使用 Megatron L文档教程开发工具CogVideoX高性能训练实战Slurm与PBS集群调度终极指南CogVideoX高性能训练实战Slurm与PBS集群调度终极指南 CogVideoX作为领先的文本与图像到视频生成模型其高性能训练需要强大的集群调度支持。人工智能大模型媒体生成预训练微调如何在 Slurm 集群上运行 Faiss 分布式 k-means 训练大量质心如何在 Slurm 集群上运行 Faiss 分布式 k means 训练大量质心 当你需要给上亿向量级的训练集聚出百万级质心例如用 IVF 索引切 100机器学习搜索引擎向量数据库上一篇5分钟掌握WeTextProcessing一站式文本标准化处理的开源利器下一篇OpenCVE部署架构详解Docker容器化部署和配置最佳实践创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考