{"id":1735,"date":"2026-06-23T11:55:15","date_gmt":"2026-06-23T03:55:15","guid":{"rendered":"https:\/\/www.izhuhn.cn\/?p=1735"},"modified":"2026-06-23T11:55:15","modified_gmt":"2026-06-23T03:55:15","slug":"dit4dit-%e6%98%87%e8%85%benpu%e9%80%82%e9%85%8d%e7%ae%80%e8%ae%b0","status":"publish","type":"post","link":"https:\/\/www.izhuhn.cn\/index.php\/2026\/06\/23\/dit4dit-%e6%98%87%e8%85%benpu%e9%80%82%e9%85%8d%e7%ae%80%e8%ae%b0\/","title":{"rendered":"DiT4DiT \u6607\u817eNPU\u9002\u914d\u7b80\u8bb0"},"content":{"rendered":"\n

DiT4DiT \u662f\u6e2f\u79d1\u5e7f\u56e2\u961f\u5728 StarVLA \u57fa\u7840\u4e0a\u63d0\u51fa\u7684\u89c6\u89c9-\u52a8\u4f5c\u6a21\u578b\uff08VAM\uff09\uff0c\u5c06\u89c6\u9891\u751f\u6210 Diffusion Transformer \u4e0e Flow Matching \u52a8\u4f5c\u9884\u6d4b\u7ed3\u5408\uff0c\u652f\u6301\u673a\u68b0\u81c2\u7075\u5de7\u64cd\u4f5c\u4e0e\u4eba\u5f62\u673a\u5668\u4eba\u5168\u8eab\u63a7\u5236\u3002\u539f\u59cb\u4ee3\u7801\u57fa\u4e8e CUDA \u751f\u6001\uff0c\u672c\u6587\u8bb0\u5f55\u5c06\u5176\u9002\u914d\u5230\u534e\u4e3a\u6607\u817e Atlas 800T A3 NPU \u7684\u8fc7\u7a0b\u3002<\/p>\n\n\n\n

\u4e00\u3001\u9879\u76ee\u5173\u8054\uff1aDiT4DiT \u4e0e StarVLA<\/h2>\n\n\n\n

\u4e24\u4e2a\u5f00\u6e90\u4ed3\u5e93\u7684\u5173\u7cfb\u5982\u4e0b\uff1a<\/p>\n\n\n\n

\n\n
<\/th>StarVLA<\/th>DiT4DiT<\/th><\/tr><\/thead>\n
\u5b9a\u4f4d<\/td>VLA \u57fa\u7840\u6846\u67b6<\/td>\u57fa\u4e8e StarVLA \u7684\u6269\u5c55<\/td><\/tr>\n
\u57fa\u5ea7\u6a21\u578b<\/td>Qwen3-VL-4B-Instruct<\/td>Cosmos-Predict2.5-2B<\/td><\/tr>\n
\u52a8\u4f5c\u9884\u6d4b<\/td>\u6807\u51c6 diffusion head<\/td>Flow-matching action head<\/td><\/tr>\n
\u89c6\u9891\u6a21\u578b<\/td>\u65e0<\/td>\u89c6\u9891 diffusion + VAE \u7f16\u89e3\u7801<\/td><\/tr>\n
\u534e\u4e3a\u9002\u914d<\/td>\u2705 \u5df2\u6709 NPU \u7248\uff08DrivingSDK\uff09<\/td>\u274c \u6682\u65e0<\/td><\/tr>\n
\u4ed3\u5e93\u5730\u5740<\/td>gitcode.com\/Ascend\/DrivingSDK<\/td>github.com\/Mondo-Robotics\/DiT4DiT<\/td><\/tr>\n<\/tbody><\/table><\/figure>\n\n\n\n

DiT4DiT \u590d\u7528\u4e86 StarVLA \u7684 dataloader\uff08LeRobot \u683c\u5f0f\uff09\u3001accelerate + DeepSpeed \u8bad\u7ec3\u6846\u67b6\u4ee5\u53ca DiT-B \u52a8\u4f5c\u5934\uff0c\u4f46\u5c06\u89c6\u89c9\u7f16\u7801\u5668\u4ece Qwen3-VL \u66ff\u6362\u4e3a Cosmos-Predict2.5-2B\uff08NVIDIA \u7684\u89c6\u9891\u751f\u6210\u6269\u6563\u6a21\u578b\uff09\uff0c\u5e76\u589e\u52a0\u4e86 Flow Matching \u52a8\u4f5c\u9884\u6d4b + \u89c6\u9891\u8f85\u52a9\u635f\u5931\u3002<\/p>\n\n\n\n

\u9002\u914d\u601d\u8def\uff1a\u53c2\u7167\u534e\u4e3a DrivingSDK \u4e2d\u5df2\u8dd1\u901a\u7684 StarVLA NPU \u7248\uff08starvla.patch\uff09\uff0c\u5bf9 DiT4DiT \u505a\u76f8\u540c\u7684 NPU \u79fb\u690d\u64cd\u4f5c\u3002<\/p>\n\n\n\n

\u4e8c\u3001\u73af\u5883\u51c6\u5907<\/h2>\n\n\n\n

2.1 \u8f6f\u4ef6\u7248\u672c<\/h3>\n\n\n\n
\n\n
\u8f6f\u4ef6<\/th>\u7248\u672c<\/th><\/tr><\/thead>\n
CANN<\/td>9.0.0<\/td><\/tr>\n
PyTorch<\/td>2.7.1<\/td><\/tr>\n
torch_npu<\/td>\u968f CANN \u9644\u8d60<\/td><\/tr>\n
DeepSpeed<\/td>0.16.9<\/td><\/tr>\n
accelerate<\/td>1.12.0<\/td><\/tr>\n
transformers<\/td>4.57.6<\/td><\/tr>\n
diffusers<\/td>git main\uff08Cosmos2.5 \u652f\u6301\uff09<\/td><\/tr>\n<\/tbody><\/table><\/figure>\n\n\n\n

2.2 \u6a21\u578b\u6743\u91cd\u4e0b\u8f7d<\/h3>\n\n\n\n
huggingface-cli download nvidia\/Cosmos-Predict2.5-2B \\\n  --revision diffusers\/base\/post-trained \\\n  --local-dir playground\/Pretrained_models\/Cosmos-Predict2.5-2B<\/code><\/pre>\n\n\n\n
2.3 LD_LIBRARY_PATH<\/h3>\n\n\n\n
conda \u73af\u5883\u4e0b torch_npu \u7684 .so \u6587\u4ef6\u4e0d\u4f1a\u81ea\u52a8\u52a0\u5165\u52a8\u6001\u5e93\u641c\u7d22\u8def\u5f84\uff0c\u9700\u5728\u542f\u52a8\u811a\u672c\u4e2d\u663e\u5f0f\u8bbe\u7f6e\uff1a<\/p>\n\n\n\n
export LD_LIBRARY_PATH=\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch\/lib:\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch_npu\/lib:$LD_LIBRARY_PATH<\/code><\/pre>\n\n\n\n
\u4e5f\u53ef\u5199\u5165 conda activate \u94a9\u5b50\u5b9e\u73b0\u81ea\u52a8\u751f\u6548\uff1a<\/p>\n\n\n\n
mkdir -p \/opt\/conda\/envs\/torch2.7.1\/etc\/conda\/activate.d\ncat > \/opt\/conda\/envs\/torch2.7.1\/etc\/conda\/activate.d\/ld_path.sh << 'EOF'\nexport LD_LIBRARY_PATH=\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch\/lib:\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch_npu\/lib:$LD_LIBRARY_PATH\nEOF<\/code><\/pre>\n\n\n\n
\u4e09\u3001\u4ee3\u7801\u9002\u914d<\/h2>\n\n\n\n
\u9002\u914d\u53ea\u9700\u4fee\u6539\u4e24\u4e2a\u6587\u4ef6\uff0c\u6838\u5fc3\u6539\u52a8\u70b9\u5982\u4e0b\uff1a<\/p>\n\n\n\n
3.1 train.py\uff08\u8bad\u7ec3\u5165\u53e3\uff09<\/h3>\n\n\n\n
    \n
  • torch_npu + patcher \u5bfc\u5165<\/strong>\uff1a\u5728 import torch.distributed<\/code> \u540e\u63d2\u5165<\/li>\n
  • autocast \u66ff\u6362<\/strong>\uff1atorch.autocast(\"cuda\")<\/code> \u2192 torch.autocast(\"npu\")<\/code><\/li>\n
  • AdamW \u878d\u5408\u4f18\u5316<\/strong>\uff1a\u52a0 fused=True<\/code><\/li>\n
  • patcher \u521d\u59cb\u5316<\/strong>\uff1a__main__<\/code> \u5165\u53e3\u5904\u8c03\u7528 default_patcher_builder.build().__enter__()<\/code><\/li>\n<\/ul>\n\n\n\n
    # === \u6539\u52a8 1: import \u533a\u57df ===\nimport torch\nimport torch.distributed as dist\n\nimport warnings\n\ntry:\n    import torch_npu\n    from torch_npu.contrib import transfer_to_npu\nexcept ImportError as e:\n    warnings.warn(f\"Failed to import torch_npu or its submodule: {e}\", ImportWarning)\n\ntry:\n    from mx_driving.patcher import default_patcher_builder\nexcept ImportError as e:\n    warnings.warn(f\"Failed to import from mx_driving.patcher: {e}\", ImportWarning)\n\nwarnings.filterwarnings(\"ignore\", category=DeprecationWarning, module=\"pandas\")\n\n# === \u6539\u52a8 2: AdamW fused=True ===\noptimizer = torch.optim.AdamW(\n    param_groups,\n    ...\n    fused=True,\n)\n\n# === \u6539\u52a8 3: autocast cuda \u2192 npu ===\n# torch.autocast(\"cuda\", dtype=torch.bfloat16)\ntorch.autocast(\"npu\", dtype=torch.bfloat16)\n\n# === \u6539\u52a8 4: __main__ \u5165\u53e3 ===\nif __name__ == \"__main__\":\n    patcher = default_patcher_builder.build()\n    patcher.__enter__()<\/code><\/pre>\n\n\n\n
    3.2 DiT4DiT.py\uff08\u6a21\u578b\u6846\u67b6\uff09<\/h3>\n\n\n\n
    \u8be5\u6587\u4ef6\u6709 4 \u5904 torch.autocast(\"cuda\")<\/code>\uff0c\u5168\u90e8\u66ff\u6362\u4e3a torch.autocast(\"npu\")<\/code>\uff1a<\/p>\n\n\n\n
      \n
    • forward() \u4e2d backbone \u524d\u5411\uff08bf16\uff09<\/li>\n
    • forward() \u4e2d action \u524d\u5411\uff08fp32\uff09<\/li>\n
    • predict_action() \u4e2d backbone \u524d\u5411\uff08bf16\uff09<\/li>\n
    • predict_action() \u4e2d action \u89e3\u7801\uff08fp32\uff09<\/li>\n<\/ul>\n\n\n\n
      3.3 \u9002\u914d\u603b\u7ed3<\/h3>\n\n\n\n
      \n\n
      \u6539\u52a8\u9879<\/th>\u539f\u56e0<\/th><\/tr><\/thead>\n
      torch_npu import + transfer_to_npu<\/td>\u66ff\u6362 torch.cuda.* \u4e3a torch.npu.*<\/td><\/tr>\n
      default_patcher_builder<\/td>mx_driving \u7684 NPU monkey patch\uff08mmcv\/torch\/numpy\uff09<\/td><\/tr>\n
      autocast(\"npu\")<\/td>NPU \u6df7\u5408\u7cbe\u5ea6\uff08\u5bf9\u5e94 CUDA \u7684 amp\uff09<\/td><\/tr>\n
      AdamW fused=True<\/td>NPU \u878d\u5408\u4f18\u5316\u5668\uff08\u51cf\u5c11 kernel launch \u5f00\u9500\uff09<\/td><\/tr>\n
      LD_LIBRARY_PATH<\/td>libc10.so\u3001libtorch_npu.so \u52a8\u6001\u5e93\u8def\u5f84<\/td><\/tr>\n<\/tbody><\/table><\/figure>\n\n\n\n
      \u6ce8\u610f\uff1aDiT4DiT \u4f7f\u7528\u7684 default_patcher_builder<\/code> API \u4e0e StarVLA \u539f\u59cb starvla.patch \u4e2d\u7684 Patcher().add(TransformersNPU)<\/code> \u4e0d\u540c\uff0c\u539f\u56e0\u662f\u65b0\u7248\u672c mx_driving \u5c06 Patcher \u4ece\u65e0\u53c2\u6784\u9020\u6539\u4e3a\u4e86 PatcherBuilder \u6a21\u5f0f\uff0cTransformersNPU<\/code> \u7c7b\u4e5f\u5df2\u79fb\u9664\uff08\u529f\u80fd\u6574\u5408\u8fdb\u4e86 default_patcher_builder\uff09\u3002<\/p>\n\n\n\n
      \u56db\u30018\u5361\u542f\u52a8\u811a\u672c<\/h2>\n\n\n\n
      \u5b8c\u6574\u7684 train_8p.sh\uff1a<\/p>\n\n\n\n
      #!\/bin\/bash\nexport WANDB_MODE=offline\nexport WANDB_OFFLINE=true\nexport TASK_QUEUE_ENABLE=2\nexport PYTORCH_NPU_ALLOC_CONF=expandable_segments:True\nexport CPU_AFFINITY_CONF=1\n#export TORCH_HCCL_ZERO_COPY=1  # A2 \u9700\u6ce8\u91ca\nexport LD_LIBRARY_PATH=\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch\/lib:\/opt\/conda\/envs\/torch2.7.1\/lib\/python3.10\/site-packages\/torch_npu\/lib:$LD_LIBRARY_PATH\n\nnum_processes=8\nmax_train_steps=150000\nper_device_batch_size=4\ndata_mix=robotwin\nbase_model=.\/playground\/Pretrained_models\/Cosmos-Predict2.5-2B\nconfig_yaml=.\/DiT4DiT\/config\/robotwin\/dit4dit_robotwin.yaml\nFramework_name=DiT4DiT\nfreeze_module_list=''\nrun_root_dir=.\/results\/Checkpoints\naction_video_freq_ratio=1\n\n# ... \u53c2\u6570\u89e3\u6790\uff08\u7565\uff09\n\naccelerate launch \\\n  --config_file DiT4DiT\/config\/deepseeds\/deepspeed_zero2.yaml \\\n  --num_processes ${num_processes} \\\n  DiT4DiT\/training\/train.py \\\n  --config_yaml ${config_yaml} \\\n  --framework.name ${Framework_name} \\\n  --framework.cosmos25.base_model ${base_model} \\\n  --datasets.vla_data.per_device_batch_size ${per_device_batch_size} \\\n  --datasets.vla_data.data_mix ${data_mix} \\\n  --trainer.max_train_steps ${max_train_steps} \\\n  --run_root_dir ${run_root_dir} \\\n  --run_id ${run_id} \\\n  2>&1 | tee ${LOG_FILE}<\/code><\/pre>\n\n\n\n
      \u4f7f\u7528\u65b9\u5f0f\uff1a<\/p>\n\n\n\n
      bash train_8p.sh --num_processes=8 --max_train_steps=100 --per_device_batch_size=1<\/code><\/pre>\n\n\n\n
      \u4e94\u3001\u8bad\u7ec3\u914d\u7f6e\uff08Robotwin\uff09<\/h2>\n\n\n\n
      \u65b0\u5efa\u914d\u7f6e\u6587\u4ef6 DiT4DiT\/config\/robotwin\/dit4dit_robotwin.yaml<\/code>\uff0c\u5173\u952e\u53c2\u6570\uff1a<\/p>\n\n\n\n
        \n
      • base_model<\/strong>: Cosmos-Predict2.5-2B\uff08\u89c6\u9891\u751f\u6210\u6269\u6563\u6a21\u578b\uff09<\/li>\n
      • data_mix<\/strong>: robotwin\uff08\u590d\u7528 StarVLA \u6570\u636e\u96c6\uff09<\/li>\n
      • action_dim\/state_dim<\/strong>: 14\uff08\u53cc\u81c2 7 \u5173\u8282 \u00d7 2\uff09<\/li>\n
      • video_delta_indices<\/strong>: [0-7]\uff088 \u5e27\u89c6\u9891\u5e8f\u5217\uff09<\/li>\n
      • training<\/strong>: joint\uff08\u52a8\u4f5c + \u89c6\u9891\u8054\u5408\u8bad\u7ec3\uff09<\/li>\n
      • future_loss_type<\/strong>: flow_matching\uff08\u6f5c\u7a7a\u95f4\u6d41\u5339\u914d\u635f\u5931\uff09<\/li>\n
      • future_action_window_size<\/strong>: 7<\/li>\n<\/ul>\n\n\n\n
        \u516d\u3001\u8e29\u5751\u8bb0\u5f55<\/h2>\n\n\n\n
        6.1 num_processes \u5fc5\u987b\u5339\u914d\u7269\u7406\u5361\u6570<\/h3>\n\n\n\n
        train_8p.sh \u540d\u5e26 8p\uff0c\u4f46\u5b9e\u9645\u4f7f\u7528 --num_processes=16<\/code> \u4f1a\u5bfc\u81f4 rank 0 \u548c rank 8 \u5206\u914d\u5230\u540c\u4e00\u5f20\u7269\u7406\u5361\uff0cHCCL \u62a5\u9519\uff1asame physical device ID0 as the rank8<\/code>\u3002\u5fc5\u987b\u4fdd\u6301 --num_processes=8<\/code>\u3002<\/p>\n\n\n\n
        6.2 mx_driving.patcher API \u53d8\u66f4<\/h3>\n\n\n\n
        \u65e7\u7248 StarVLA starvla.patch \u4f7f\u7528 Patcher().add(TransformersNPU).apply()<\/code>\uff0c\u65b0\u7248 mx_driving \u6539\u4e3a\u4e86 default_patcher_builder.build().__enter__()<\/code>\u3002TransformersNPU<\/code> \u7c7b\u5df2\u79fb\u9664\uff0c\u529f\u80fd\u6574\u5408\u8fdb\u4e86 default_patcher_builder \u9884\u8bbe\u7684 mmcv\/torch\/numpy\/mmdet \u5168\u5957 patch\u3002<\/p>\n\n\n\n
        6.3 libc10.so \/ libtorch_npu.so \u672a\u81ea\u52a8\u53d1\u73b0<\/h3>\n\n\n\n
        \u4ea4\u4e92\u5f0f python shell \u4e2d import mx_driving \u4f1a\u62a5 libc10.so: cannot open shared object file<\/code>\uff0c\u9700\u624b\u52a8\u8bbe\u7f6e LD_LIBRARY_PATH\u3002\u8bad\u7ec3\u811a\u672c\u4e2d torch_npu \u901a\u8fc7 transfer_to_npu<\/code> \u5728\u8fd0\u884c\u65f6\u4f1a\u6b63\u786e\u8bbe\u7f6e\u8def\u5f84\uff0c\u4f46\u6392\u67e5\u95ee\u9898\u65f6\u9700\u6ce8\u610f\u3002<\/p>\n\n\n\n
        6.4 Cosmos-Predict2.5-2B \u663e\u5b58\u5360\u7528<\/h3>\n\n\n\n
        Cosmos2.5 \u5305\u542b VAE + Transformer + Text Encoder \u4e09\u90e8\u5206\uff0c\u5b8c\u6574\u52a0\u8f7d\u7ea6\u9700 10GB+ \u663e\u5b58\u3002\u914d\u5408 DeepSpeed ZeRO-2 \u53ef\u5c06\u4f18\u5316\u5668\u72b6\u6001\u5206\u7247\u5230 8 \u5f20\u5361\uff0c\u4f46\u9996\u5361\u4ecd\u9700\u627f\u8f7d\u5b8c\u6574\u6a21\u578b\u6743\u91cd\uff08\u7ea6 4.5B \u53c2\u6570\uff09\u3002\u5b9e\u9645\u8bad\u7ec3\u9700\u6839\u636e\u5361\u663e\u5b58\uff08A3 \u5355\u5361\u7ea6 64GB HBM\uff09\u8c03\u6574 per_device_batch_size\u3002<\/p>\n\n\n\n
        \u7b80\u8bb0\u3002<\/p>\n","protected":false},"excerpt":{"rendered":"
        DiT4DiT \u662f\u6e2f\u79d1\u5e7f\u56e2\u961f\u5728 StarVLA \u57fa\u7840\u4e0a\u63d0\u51fa\u7684\u89c6\u89c9-\u52a8\u4f5c\u6a21\u578b\uff08VAM\uff09\uff0c\u5c06\u89c6\u9891\u751f\u6210 Diffusion Transf …<\/p>\n","protected":false},"author":1,"featured_media":0,"comment_status":"closed","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[],"class_list":["post-1735","post","type-post","status-publish","format-standard","hentry","category-uncategorized"],"_links":{"self":[{"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/posts\/1735","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/comments?post=1735"}],"version-history":[{"count":1,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/posts\/1735\/revisions"}],"predecessor-version":[{"id":1736,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/posts\/1735\/revisions\/1736"}],"wp:attachment":[{"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/media?parent=1735"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/categories?post=1735"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.izhuhn.cn\/index.php\/wp-json\/wp\/v2\/tags?post=1735"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}