torchtitan-npu Checkpoint 完整使用指南DCP 断点续训、Hugging Face 权重加载与 seed checkpoint 实战【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu本指南系统讲解如何在torchtitan-npuAscend Extension for torchtitan中使用 TorchTitan 的 checkpoint 机制包括基于 DCPDistributed Checkpoint的分布式训练状态保存/恢复、基于 Hugging Face safetensors 的模型权重加载与导出、以及面向多卡重分片场景的 seed checkpoint 创建。读完本文你将掌握 checkpoint 的目录格式、全部常用配置项含默认值与约束、可直接复制的命令行与 Python 配置写法并理解 NPU 侧对 checkpoint 的扩展实现同步写入、SHA-256 完整性校验、EMA 状态管理与底层源码依据。一、核心机制DCP 与 Hugging Face safetensors 的分工torchtitan-npu的 checkpoint 能力建立在 TorchTitan 上游的torchtitan.components.checkpointer.CheckpointManager之上torchtitan v0.3.0 起torchtitan.components.checkpoint仅为兼容再导出模块。两种格式各司其职务必区分DCPDistributed Checkpoint用于保存和加载完整分布式训练状态涵盖模型、优化器、学习率调度器、dataloader 与训练步数适合断点续训也支持改变并行切分后的重分片加载。Hugging Face safetensors只用于模型权重的加载或保存不能加载优化器和训练步数不能作为完整训练断点使用。需要注意的是本指南不覆盖权重转换工具。需要将 Hugging Face 权重转换为 TorchTitan 格式时请使用上游仓库提供的转换工具本仓库 patches/torchtitan/scripts/checkpoint_conversion/convert_to_hf.py 中也包含相关转换实现的参考其中ParallelFileSystemReader提供了并行读取 DCP 分片的能力其行为在 tests/unit_tests/ema/test_ema_initial_load.py 中有单测覆盖。二、当前配置入口与默认值训练配置由各模型的配置注册表提供checkpoint 部分直接使用CheckpointManager.Config构建。当前仓库中涉及 checkpoint 默认值的注册表如下模块配置入口checkpoint 默认值torchtitan_npu/models/deepseek_v4/config_registry.pydeepseek_v4_debugmodel、deepseek_v4_flash、deepseek_v4_proenableFalseinterval100torchtitan_npu/models/deepseek_v3_2/config_registry.pydeepseek_v3_2_debugmodelenableFalseinterval10last_save_model_onlyFalse从源码看deepseek_v4系列在 config_registry.py 中构造CheckpointManager.Config(enableFalse, interval100)而deepseek_v3_2_debugmodel在 config_registry.py 中构造CheckpointManager.Config(interval10, last_save_model_onlyFalse)。关键行为enableFalse时既不会保存 checkpoint也不会执行 checkpoint 加载。因此训练命令需要显式传入--checkpoint.enable或在 Python 配置中设置checkpoint.enable True。三、checkpoint 格式与目录结构3.1 DCP 目录结构DCP 是 PyTorch Distributed Checkpoint 格式。启用checkpoint.foldercheckpoint时目录结构通常为dump_folder/checkpoint/step-step/ ├── .metadata └── *.distcp其中dump_folder来自--dump-folder命令行参数即{dump_folder}/{folder}拼接而成folder是 checkpoint 子目录名step-step目录名中的step为训练步数DCP 加载会检查.metadata并按--checkpoint.load-step选择指定步数默认值-1表示选择最新步数。3.2 Hugging Face safetensors 形态Hugging Face 权重通常包含model.safetensors或model.safetensors.index.json分片索引使用前还应准备对应的 tokenizer 与config.json等配置文件。该格式只表示模型权重不包含优化器、学习率调度器和训练步数。四、常用配置项全解配置类型来自上游torchtitan.config.Checkpoint。Python 配置使用下划线命名如load_step命令行使用 tyro 的连字符命名如--checkpoint.load-step。下表列出全部常用配置项配置项作用enable启用 checkpoint 保存和加载。默认False。foldercheckpoint 子目录名默认checkpoint最终路径为{dump_folder}/{folder}。interval保存 DCP checkpoint 的步数间隔。load_step要加载的步数-1表示最新 checkpoint。initial_load_path当前输出目录没有 checkpoint 时使用的初始 checkpoint 路径必须是绝对路径或远程 URI。initial_load_model_only初始加载是否只加载模型权重默认True。设为False才会尝试加载完整训练状态。initial_load_in_hf将初始路径按 Hugging Face safetensors 读取。HF 加载只能加载模型权重。initial_load_in_hf_quantized从 HF 量化权重加载使用前必须启用initial_load_in_hf。last_save_model_only最后一步是否只保存模型权重默认True。设为False才保存完整训练状态。last_save_in_hf最后一步是否以 Hugging Face safetensors 格式保存。必须同时使用模型权重保存模式。export_dtype保存时模型权重导出的 dtype可用float16、bfloat16、float32。async_mode保存方式disabled、async或async_with_pinned_mem。keep_latest_k保留最近的 checkpoint 数量0表示全部保留不能设置为1。exclude_from_loading从 DCP 加载时排除状态例如optimizer,lr_scheduler,dataloader。enable_first_step_checkpoint是否在第一个训练 step 后立即保存一次 checkpoint。create_seed_checkpoint不应用并行切分创建可供后续任务重分片加载的 seed checkpoint。load_only只加载、不保存 checkpoint适合验证或调试。命令行书写规则布尔字段使用反向选项关闭例如--checkpoint.no-enable不要写成--checkpoint.enable false。列表字段使用英文逗号分隔不要用空格拆成多个 token例如--checkpoint.exclude-from-loading optimizer,lr_scheduler,dataloader。五、保存 DCP checkpoint5.1 命令行方式以下命令使用当前脚本默认的deepseek_v3_debugmodel配置MODULEtorchtitan.models.deepseek_v3、CONFIGdeepseek_v3_debugmodel见 scripts/run_train.sh。运行前需准备 Ascend/CANN 环境脚本会自动 source CANN 的set_env.sh并默认将 Inductor 后端设置为ascendcNGPU1 \ bash scripts/run_train.sh \ --hf-assets-path tests/assets/deepseek_v3 \ --dump-folder ./outputs/dsv3_checkpoint \ --checkpoint.enable \ --checkpoint.folder checkpoint \ --checkpoint.interval 100 \ --training.steps 200训练过程中会按interval100生成./outputs/dsv3_checkpoint/checkpoint/step-100/ ./outputs/dsv3_checkpoint/checkpoint/step-200/5.2 Python 配置方式也可以在配置注册表中直接设置CheckpointManager.Config的字段与命令行一一对应from torchtitan.components.checkpointer import CheckpointManager checkpoint CheckpointManager.Config( enableTrue, foldercheckpoint, interval100, keep_latest_k5, async_modedisabled, )六、加载 DCP checkpoint6.1 自动加载最新 checkpoint使用相同的--dump-folder和--checkpoint.folder重新启动并启用 checkpointNGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_checkpoint \ --checkpoint.enable \ --checkpoint.folder checkpoint当输出目录中存在可用的step-*checkpoint 时框架会默认加载最新一步对应load_step-1。6.2 指定加载步数bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_checkpoint \ --checkpoint.enable \ --checkpoint.load-step 1006.3 从其他目录初始化使用新的输出目录并通过--checkpoint.initial-load-path指向旧 checkpoint 的完整 step 目录。该参数必须使用绝对路径或远程 URI相对路径会触发上游的ValueErrorNGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_new_job \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100默认只加载模型权重。若要加载优化器、学习率调度器和训练状态应显式关闭initial_load_model_onlybash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_new_job \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100 \ --checkpoint.no-initial-load-model-only优先级与冲突处理如果{dump_folder}/{checkpoint.folder}已经存在可用 checkpoint框架会优先从该目录加载并忽略initial_load_path。因此从新权重启动实验时应使用新的dump-folder或清理旧的 checkpoint 目录。如只需要模型和部分训练状态可以排除不需要的键bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_model_only \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100 \ --checkpoint.exclude-from-loading optimizer,lr_scheduler,dataloader七、加载 Hugging Face 权重从 HF safetensors 初始化时需要启用 checkpoint、设置--checkpoint.initial-load-in-hf并确保--checkpoint.initial-load-model-only保持为TrueHF 加载天然只支持模型权重NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_from_hf \ --checkpoint.enable \ --checkpoint.initial-load-in-hf \ --checkpoint.initial-load-path /absolute/path/to/checkpoint/DeepSeek-V4路径解析规则如果不传--checkpoint.initial-load-pathcheckpoint 管理器会尝试使用模型配置中的hf_assets_pathinitial_load_path的优先级高于hf_assets_path。HF 权重到本地模型参数的映射由各模型的 state-dict adapter 完成。以 DeepSeek-V4 为例state_dict_adapter.py 中的from_hf_map定义了完整的键映射关系例如词嵌入与输出头embed.weight↔tok_embeddings.weight、head.weight↔lm_head.weight注意力layers.{}.attn.wq_a.weight↔layers.{}.attention.wq_a.weight、layers.{}.attn.wo_b.weight↔layers.{}.attention.wo_b.weightMoE 路由专家layers.{}.ffn.experts.{}.w1.weight↔layers.{}.moe.routed_experts.inner_experts.w1_EFD反向导出时还会按专家数拆分/拼接权重涉及grouped_expert_weight_placements等 DTensor 信息见 state_dict_adapter.py压缩器compressor与索引器indexer仅在对应层compress_ratios ! 1或 4时按层追加映射MTP 层本地mtp_layers.{depth}.*命名空间与 HF 的mtp.{depth}.*双向换算见 state_dict_adapter.py。如果 HF checkpoint 中存在超出模型配置的 MTP stage 数量adapter 会抛出ValueError而不是静默丢弃。八、保存 Hugging Face 权重训练最后一步可以直接保存 HF safetensors。该模式只保存模型权重不能作为完整训练断点使用NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_hf \ --checkpoint.enable \ --checkpoint.last-save-in-hf \ --checkpoint.export-dtype bfloat16 \ --training.steps 100使用前提last_save_in_hf需要模型提供 state-dict adapter必须保持last_save_model_onlyTrue从源码结构看deepseek_v3_2与deepseek_v4模型均提供对应的 state-dict adapterdeepseek_v3_2/state_dict_adapter.py、deepseek_v4/state_dict_adapter.py但具体可用性仍取决于模型配置与权重格式export_dtype决定导出权重精度可选float16、bfloat16、float32。九、创建 seed checkpointseed checkpoint 用于先创建未应用并行切分的模型状态再由多卡任务通过 DCP 重分片加载。创建时使用单卡并将各并行度设为1NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_seed \ --checkpoint.enable \ --checkpoint.create-seed-checkpoint \ --parallelism.data-parallel-replicate-degree 1 \ --parallelism.data-parallel-shard-degree 1 \ --parallelism.tensor-parallel-degree 1 \ --parallelism.pipeline-parallel-degree 1 \ --parallelism.context-parallel-degree 1 \ --parallelism.expert-parallel-degree 1生成的step-0目录可以直接作为后续任务的checkpoint.initial_load_path。由于 seed checkpoint 的模型状态不依赖任何并行切分后续任务可以按自身的并行方案例如 FSDP TP EP CP在加载时通过 DCP 机制重分片这正是 DCP 相对朴素权重 dump 的核心优势。十、NPU 侧 checkpoint 扩展源码级深度解读除了上游CheckpointManager的能力torchtitan-npu在 NPU 场景下提供了三层扩展实现可从源码与测试中逐一印证10.1 NPU CheckpointManager 扩展同步写入 完整性校验tests/unit_tests/extensions/components/test_checkpoint.py 通过 stub 上游模块的方式对 torchtitan_npu/extensions/components/checkpoint.py 做了纯 CPU 的单测覆盖其行为可归纳为扩展配置CheckpointManager.Config新增extensions: CheckpointExtensions字段其中verify_hash_manifest默认False见 checkpoint.py 与测试test_config_defaults_hash_manifest_verification_to_disabled。同步写入路径当async_mode AsyncMode.DISABLED且非 HF 导出时不走上游的异步 stager而是直接以dcp.FileSystemWriter(checkpoint_id, per_thread_copy_ahead0)同步落盘并在enable_garbage_collection时触发一次 GC见 checkpoint.py测试test_checkpoint_manager_uses_synchronous_writer。SHA-256 manifest 校验开启verify_hash_manifest后保存流程为标记 pending → 保存 checkpoint → 写入 manifest加载流程为校验 manifest → 执行上游加载。保存失败时保持 pending 状态且不写 manifest见 checkpoint.py测试test_save_marks_pending_before_save_and_writes_manifest_afterward、test_save_failure_keeps_pending_and_skips_manifest_write、test_dcp_load_verifies_manifest_before_upstream_load。manifest 中记录 checkpoint 目录内每个文件的 SHA-256 哈希哈希不匹配会抛出CheckpointManifestError并拒绝加载若 checkpoint 保存时未开启该校验manifest 缺失加载会静默放行。文件 I/O 仅由 rank 0 执行校验结论广播到所有 rank。10.2 EMA 状态与 checkpoint 联动torchtitan_npu/patches/torchtitan/components/checkpoint.py 从上游 TorchTitan PR #3985 移植了 EMA-aware checkpoint 支持上游合入后应删除该模块EMACheckpointManager继承上游CheckpointManager在Config中追加ema_weights字段构造时把 EMA 伪优化器注册进 checkpoint 状态表键为EMA_OPTIMIZER ema_optimizer加载时会检查 checkpoint 中是否包含 EMA 前缀缺失时自动从模型权重重建reseedEMA 影子参数保证只加载模型权重的路径不会留下陈旧的 EMA 状态该补丁同时替换torchtitan.components.checkpointer、torchtitan.components.checkpointer.dcp、torchtitan.components.checkpoint三个命名空间中的CheckpointManager见 checkpoint.py。对应的回归测试见 tests/unit_tests/ema/test_ema_initial_load.py覆盖了 DCP 与 HF 两种仅模型初始加载后 EMA 自动 reseedtest_native_model_only_load_reseeds_ema、test_hf_model_only_load_reseeds_ema、EMA 状态 roundtrip 与冷启动test_ema_state_roundtrip_and_cold_start等场景。EMA 本体实现在 patches/torchtitan/components/ema.py_EMAParamOptimizer为每个参数维护state[p][ema_params]影子权重并复用OptimizersContainer的 FQN 扁平化 DCP state-dict 机制天然支持重分片加载非法 EMA 配置如负 decay、decay1、零步间隔会在构建时快速失败ValueError。十一、常见限制与验证范围需要可用的torch、torch_npu、TorchTitan、CANN 和 Ascend NPU 环境本文命令的完整执行依赖该环境就绪scripts/run_train.sh 会自动探测并 source CANN 环境变量。DCP 适用于完整训练状态加载HF safetensors 只适用于模型权重二者不可混用。文档中的模型入口、配置字段和命令来自源码与脚本的静态核对多卡训练、权重格式兼容性和 NPU 性能需要在实际环境中验证。keep_latest_k不能设置为1initial_load_in_hf_quantized必须与initial_load_in_hf搭配使用last_save_in_hf必须保持模型权重保存模式。十二、相关文档快速上手安装指南训练启动脚本DeepSeek-V4 checkpoint 配置NPU checkpoint 管理器扩展EMA-aware checkpoint 补丁checkpoint 单测CPUEMA 初始加载回归测试【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
torchtitan-npu Checkpoint 完整使用指南DCP 断点续训、Hugging Face 权重加载与 seed checkpoint 实战【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu本指南系统讲解如何在torchtitan-npuAscend Extension for torchtitan中使用 TorchTitan 的 checkpoint 机制包括基于 DCPDistributed Checkpoint的分布式训练状态保存/恢复、基于 Hugging Face safetensors 的模型权重加载与导出、以及面向多卡重分片场景的 seed checkpoint 创建。读完本文你将掌握 checkpoint 的目录格式、全部常用配置项含默认值与约束、可直接复制的命令行与 Python 配置写法并理解 NPU 侧对 checkpoint 的扩展实现同步写入、SHA-256 完整性校验、EMA 状态管理与底层源码依据。一、核心机制DCP 与 Hugging Face safetensors 的分工torchtitan-npu的 checkpoint 能力建立在 TorchTitan 上游的torchtitan.components.checkpointer.CheckpointManager之上torchtitan v0.3.0 起torchtitan.components.checkpoint仅为兼容再导出模块。两种格式各司其职务必区分DCPDistributed Checkpoint用于保存和加载完整分布式训练状态涵盖模型、优化器、学习率调度器、dataloader 与训练步数适合断点续训也支持改变并行切分后的重分片加载。Hugging Face safetensors只用于模型权重的加载或保存不能加载优化器和训练步数不能作为完整训练断点使用。需要注意的是本指南不覆盖权重转换工具。需要将 Hugging Face 权重转换为 TorchTitan 格式时请使用上游仓库提供的转换工具本仓库 patches/torchtitan/scripts/checkpoint_conversion/convert_to_hf.py 中也包含相关转换实现的参考其中ParallelFileSystemReader提供了并行读取 DCP 分片的能力其行为在 tests/unit_tests/ema/test_ema_initial_load.py 中有单测覆盖。二、当前配置入口与默认值训练配置由各模型的配置注册表提供checkpoint 部分直接使用CheckpointManager.Config构建。当前仓库中涉及 checkpoint 默认值的注册表如下模块配置入口checkpoint 默认值torchtitan_npu/models/deepseek_v4/config_registry.pydeepseek_v4_debugmodel、deepseek_v4_flash、deepseek_v4_proenableFalseinterval100torchtitan_npu/models/deepseek_v3_2/config_registry.pydeepseek_v3_2_debugmodelenableFalseinterval10last_save_model_onlyFalse从源码看deepseek_v4系列在 config_registry.py 中构造CheckpointManager.Config(enableFalse, interval100)而deepseek_v3_2_debugmodel在 config_registry.py 中构造CheckpointManager.Config(interval10, last_save_model_onlyFalse)。关键行为enableFalse时既不会保存 checkpoint也不会执行 checkpoint 加载。因此训练命令需要显式传入--checkpoint.enable或在 Python 配置中设置checkpoint.enable True。三、checkpoint 格式与目录结构3.1 DCP 目录结构DCP 是 PyTorch Distributed Checkpoint 格式。启用checkpoint.foldercheckpoint时目录结构通常为dump_folder/checkpoint/step-step/ ├── .metadata └── *.distcp其中dump_folder来自--dump-folder命令行参数即{dump_folder}/{folder}拼接而成folder是 checkpoint 子目录名step-step目录名中的step为训练步数DCP 加载会检查.metadata并按--checkpoint.load-step选择指定步数默认值-1表示选择最新步数。3.2 Hugging Face safetensors 形态Hugging Face 权重通常包含model.safetensors或model.safetensors.index.json分片索引使用前还应准备对应的 tokenizer 与config.json等配置文件。该格式只表示模型权重不包含优化器、学习率调度器和训练步数。四、常用配置项全解配置类型来自上游torchtitan.config.Checkpoint。Python 配置使用下划线命名如load_step命令行使用 tyro 的连字符命名如--checkpoint.load-step。下表列出全部常用配置项配置项作用enable启用 checkpoint 保存和加载。默认False。foldercheckpoint 子目录名默认checkpoint最终路径为{dump_folder}/{folder}。interval保存 DCP checkpoint 的步数间隔。load_step要加载的步数-1表示最新 checkpoint。initial_load_path当前输出目录没有 checkpoint 时使用的初始 checkpoint 路径必须是绝对路径或远程 URI。initial_load_model_only初始加载是否只加载模型权重默认True。设为False才会尝试加载完整训练状态。initial_load_in_hf将初始路径按 Hugging Face safetensors 读取。HF 加载只能加载模型权重。initial_load_in_hf_quantized从 HF 量化权重加载使用前必须启用initial_load_in_hf。last_save_model_only最后一步是否只保存模型权重默认True。设为False才保存完整训练状态。last_save_in_hf最后一步是否以 Hugging Face safetensors 格式保存。必须同时使用模型权重保存模式。export_dtype保存时模型权重导出的 dtype可用float16、bfloat16、float32。async_mode保存方式disabled、async或async_with_pinned_mem。keep_latest_k保留最近的 checkpoint 数量0表示全部保留不能设置为1。exclude_from_loading从 DCP 加载时排除状态例如optimizer,lr_scheduler,dataloader。enable_first_step_checkpoint是否在第一个训练 step 后立即保存一次 checkpoint。create_seed_checkpoint不应用并行切分创建可供后续任务重分片加载的 seed checkpoint。load_only只加载、不保存 checkpoint适合验证或调试。命令行书写规则布尔字段使用反向选项关闭例如--checkpoint.no-enable不要写成--checkpoint.enable false。列表字段使用英文逗号分隔不要用空格拆成多个 token例如--checkpoint.exclude-from-loading optimizer,lr_scheduler,dataloader。五、保存 DCP checkpoint5.1 命令行方式以下命令使用当前脚本默认的deepseek_v3_debugmodel配置MODULEtorchtitan.models.deepseek_v3、CONFIGdeepseek_v3_debugmodel见 scripts/run_train.sh。运行前需准备 Ascend/CANN 环境脚本会自动 source CANN 的set_env.sh并默认将 Inductor 后端设置为ascendcNGPU1 \ bash scripts/run_train.sh \ --hf-assets-path tests/assets/deepseek_v3 \ --dump-folder ./outputs/dsv3_checkpoint \ --checkpoint.enable \ --checkpoint.folder checkpoint \ --checkpoint.interval 100 \ --training.steps 200训练过程中会按interval100生成./outputs/dsv3_checkpoint/checkpoint/step-100/ ./outputs/dsv3_checkpoint/checkpoint/step-200/5.2 Python 配置方式也可以在配置注册表中直接设置CheckpointManager.Config的字段与命令行一一对应from torchtitan.components.checkpointer import CheckpointManager checkpoint CheckpointManager.Config( enableTrue, foldercheckpoint, interval100, keep_latest_k5, async_modedisabled, )六、加载 DCP checkpoint6.1 自动加载最新 checkpoint使用相同的--dump-folder和--checkpoint.folder重新启动并启用 checkpointNGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_checkpoint \ --checkpoint.enable \ --checkpoint.folder checkpoint当输出目录中存在可用的step-*checkpoint 时框架会默认加载最新一步对应load_step-1。6.2 指定加载步数bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_checkpoint \ --checkpoint.enable \ --checkpoint.load-step 1006.3 从其他目录初始化使用新的输出目录并通过--checkpoint.initial-load-path指向旧 checkpoint 的完整 step 目录。该参数必须使用绝对路径或远程 URI相对路径会触发上游的ValueErrorNGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_new_job \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100默认只加载模型权重。若要加载优化器、学习率调度器和训练状态应显式关闭initial_load_model_onlybash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_new_job \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100 \ --checkpoint.no-initial-load-model-only优先级与冲突处理如果{dump_folder}/{checkpoint.folder}已经存在可用 checkpoint框架会优先从该目录加载并忽略initial_load_path。因此从新权重启动实验时应使用新的dump-folder或清理旧的 checkpoint 目录。如只需要模型和部分训练状态可以排除不需要的键bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_model_only \ --checkpoint.enable \ --checkpoint.initial-load-path /absolute/path/to/dsv4_checkpoint/checkpoint/step-100 \ --checkpoint.exclude-from-loading optimizer,lr_scheduler,dataloader七、加载 Hugging Face 权重从 HF safetensors 初始化时需要启用 checkpoint、设置--checkpoint.initial-load-in-hf并确保--checkpoint.initial-load-model-only保持为TrueHF 加载天然只支持模型权重NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_from_hf \ --checkpoint.enable \ --checkpoint.initial-load-in-hf \ --checkpoint.initial-load-path /absolute/path/to/checkpoint/DeepSeek-V4路径解析规则如果不传--checkpoint.initial-load-pathcheckpoint 管理器会尝试使用模型配置中的hf_assets_pathinitial_load_path的优先级高于hf_assets_path。HF 权重到本地模型参数的映射由各模型的 state-dict adapter 完成。以 DeepSeek-V4 为例state_dict_adapter.py 中的from_hf_map定义了完整的键映射关系例如词嵌入与输出头embed.weight↔tok_embeddings.weight、head.weight↔lm_head.weight注意力layers.{}.attn.wq_a.weight↔layers.{}.attention.wq_a.weight、layers.{}.attn.wo_b.weight↔layers.{}.attention.wo_b.weightMoE 路由专家layers.{}.ffn.experts.{}.w1.weight↔layers.{}.moe.routed_experts.inner_experts.w1_EFD反向导出时还会按专家数拆分/拼接权重涉及grouped_expert_weight_placements等 DTensor 信息见 state_dict_adapter.py压缩器compressor与索引器indexer仅在对应层compress_ratios ! 1或 4时按层追加映射MTP 层本地mtp_layers.{depth}.*命名空间与 HF 的mtp.{depth}.*双向换算见 state_dict_adapter.py。如果 HF checkpoint 中存在超出模型配置的 MTP stage 数量adapter 会抛出ValueError而不是静默丢弃。八、保存 Hugging Face 权重训练最后一步可以直接保存 HF safetensors。该模式只保存模型权重不能作为完整训练断点使用NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_hf \ --checkpoint.enable \ --checkpoint.last-save-in-hf \ --checkpoint.export-dtype bfloat16 \ --training.steps 100使用前提last_save_in_hf需要模型提供 state-dict adapter必须保持last_save_model_onlyTrue从源码结构看deepseek_v3_2与deepseek_v4模型均提供对应的 state-dict adapterdeepseek_v3_2/state_dict_adapter.py、deepseek_v4/state_dict_adapter.py但具体可用性仍取决于模型配置与权重格式export_dtype决定导出权重精度可选float16、bfloat16、float32。九、创建 seed checkpointseed checkpoint 用于先创建未应用并行切分的模型状态再由多卡任务通过 DCP 重分片加载。创建时使用单卡并将各并行度设为1NGPU1 \ HF_ASSETS_PATH/path/to/dsv4_tokenizer \ bash scripts/run_train.sh \ --dump-folder ./outputs/dsv4_seed \ --checkpoint.enable \ --checkpoint.create-seed-checkpoint \ --parallelism.data-parallel-replicate-degree 1 \ --parallelism.data-parallel-shard-degree 1 \ --parallelism.tensor-parallel-degree 1 \ --parallelism.pipeline-parallel-degree 1 \ --parallelism.context-parallel-degree 1 \ --parallelism.expert-parallel-degree 1生成的step-0目录可以直接作为后续任务的checkpoint.initial_load_path。由于 seed checkpoint 的模型状态不依赖任何并行切分后续任务可以按自身的并行方案例如 FSDP TP EP CP在加载时通过 DCP 机制重分片这正是 DCP 相对朴素权重 dump 的核心优势。十、NPU 侧 checkpoint 扩展源码级深度解读除了上游CheckpointManager的能力torchtitan-npu在 NPU 场景下提供了三层扩展实现可从源码与测试中逐一印证10.1 NPU CheckpointManager 扩展同步写入 完整性校验tests/unit_tests/extensions/components/test_checkpoint.py 通过 stub 上游模块的方式对 torchtitan_npu/extensions/components/checkpoint.py 做了纯 CPU 的单测覆盖其行为可归纳为扩展配置CheckpointManager.Config新增extensions: CheckpointExtensions字段其中verify_hash_manifest默认False见 checkpoint.py 与测试test_config_defaults_hash_manifest_verification_to_disabled。同步写入路径当async_mode AsyncMode.DISABLED且非 HF 导出时不走上游的异步 stager而是直接以dcp.FileSystemWriter(checkpoint_id, per_thread_copy_ahead0)同步落盘并在enable_garbage_collection时触发一次 GC见 checkpoint.py测试test_checkpoint_manager_uses_synchronous_writer。SHA-256 manifest 校验开启verify_hash_manifest后保存流程为标记 pending → 保存 checkpoint → 写入 manifest加载流程为校验 manifest → 执行上游加载。保存失败时保持 pending 状态且不写 manifest见 checkpoint.py测试test_save_marks_pending_before_save_and_writes_manifest_afterward、test_save_failure_keeps_pending_and_skips_manifest_write、test_dcp_load_verifies_manifest_before_upstream_load。manifest 中记录 checkpoint 目录内每个文件的 SHA-256 哈希哈希不匹配会抛出CheckpointManifestError并拒绝加载若 checkpoint 保存时未开启该校验manifest 缺失加载会静默放行。文件 I/O 仅由 rank 0 执行校验结论广播到所有 rank。10.2 EMA 状态与 checkpoint 联动torchtitan_npu/patches/torchtitan/components/checkpoint.py 从上游 TorchTitan PR #3985 移植了 EMA-aware checkpoint 支持上游合入后应删除该模块EMACheckpointManager继承上游CheckpointManager在Config中追加ema_weights字段构造时把 EMA 伪优化器注册进 checkpoint 状态表键为EMA_OPTIMIZER ema_optimizer加载时会检查 checkpoint 中是否包含 EMA 前缀缺失时自动从模型权重重建reseedEMA 影子参数保证只加载模型权重的路径不会留下陈旧的 EMA 状态该补丁同时替换torchtitan.components.checkpointer、torchtitan.components.checkpointer.dcp、torchtitan.components.checkpoint三个命名空间中的CheckpointManager见 checkpoint.py。对应的回归测试见 tests/unit_tests/ema/test_ema_initial_load.py覆盖了 DCP 与 HF 两种仅模型初始加载后 EMA 自动 reseedtest_native_model_only_load_reseeds_ema、test_hf_model_only_load_reseeds_ema、EMA 状态 roundtrip 与冷启动test_ema_state_roundtrip_and_cold_start等场景。EMA 本体实现在 patches/torchtitan/components/ema.py_EMAParamOptimizer为每个参数维护state[p][ema_params]影子权重并复用OptimizersContainer的 FQN 扁平化 DCP state-dict 机制天然支持重分片加载非法 EMA 配置如负 decay、decay1、零步间隔会在构建时快速失败ValueError。十一、常见限制与验证范围需要可用的torch、torch_npu、TorchTitan、CANN 和 Ascend NPU 环境本文命令的完整执行依赖该环境就绪scripts/run_train.sh 会自动探测并 source CANN 环境变量。DCP 适用于完整训练状态加载HF safetensors 只适用于模型权重二者不可混用。文档中的模型入口、配置字段和命令来自源码与脚本的静态核对多卡训练、权重格式兼容性和 NPU 性能需要在实际环境中验证。keep_latest_k不能设置为1initial_load_in_hf_quantized必须与initial_load_in_hf搭配使用last_save_in_hf必须保持模型权重保存模式。十二、相关文档快速上手安装指南训练启动脚本DeepSeek-V4 checkpoint 配置NPU checkpoint 管理器扩展EMA-aware checkpoint 补丁checkpoint 单测CPUEMA 初始加载回归测试【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
