Repository navigation
Feature/cosmos3 patch #19
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
iMountTai
wants to merge
9
commits into
HYGON-AI:main
Choose a base branch
from
iMountTai:feature/cosmos3-patch
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+9,691
−3
Open
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
9b4453c
feat(cosmos3): add cosmos3 adaptation and optimization
iMountTai bb302a6
Merge remote-tracking branch 'origin/main' into main
iMountTai 2d9e7fd
feat(cosmos3): integrate source patches independently
iMountTai 97db666
refactor(cosmos3): migrate patch to model_examples and streamline com…
iMountTai 4fb508b
refactor(cosmos3): clean up residual bootstrap and toml dependencies …
iMountTai af3d438
Merge branch 'main' into feature/cosmos3-patch
iMountTai 95a3044
Update README.md
iMountTai ec89451
update cosmos-framework.patch
iMountTai 298d311
docs: add Cosmos3 to model support list
iMountTai File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1 +1,3 @@ | ||
| # Preserve Git patch bytes, including whitespace in diff context. | ||
| /model_examples/Cosmos3/*.patch -text -whitespace | ||
| *.sh text eol=lf |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,207 @@ | ||
| # Cosmos3 运行指南 | ||
|
|
||
| 在 `cosmos-framework` 根目录应用 Cosmos3 补丁，再直接启动训练。以下示例使用 8 张 HCU，训练 70 步。 | ||
|
|
||
| ## 环境准备 | ||
|
|
||
| ### 拉取镜像及创建容器 | ||
| ``` | ||
| docker pull harbor.sourcefind.cn:5443/hcu/admin/base/custom:dit-vla-wm-base-ubuntu22.04-dtk26.04-torch2.10-py3.10-20260929 | ||
|
|
||
| docker run -it --network=host --name=cosmos_hygon --privileged --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size=512G --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -u root --ulimit stack=-1:-1 --ulimit memlock=-1:-1 -v /opt/hyhal:/opt/hyhal:ro harbor.sourcefind.cn:5443/hcu/admin/base/custom:dit-vla-wm-base-ubuntu22.04-dtk26.04-torch2.10-py3.10-20260929 | ||
| ``` | ||
| ### 准备代码 | ||
| ```bash | ||
| git clone --branch feature/cosmos3-patch https://github.com/iMountTai/TurboPhysAI.git | ||
| cd TurboPhysAI | ||
| cd .. | ||
| git clone https://github.com/NVIDIA/cosmos-framework.git | ||
| cd cosmos-framework | ||
| git checkout 9726697a | ||
| ``` | ||
|
|
||
| ### 应用补丁 | ||
|
|
||
| 在 Cosmos 仓库根目录应用补丁；以下假设 TurboPhysAI 和 cosmos-framework 是同级相邻目录： | ||
|
|
||
| ```bash | ||
| git apply ../TurboPhysAI/model_examples/Cosmos3/cosmos-framework.patch | ||
| ``` | ||
|
|
||
| 如需撤销补丁： | ||
|
|
||
| ```bash | ||
| git apply -R ../TurboPhysAI/model_examples/Cosmos3/cosmos-framework.patch | ||
| ``` | ||
|
|
||
| ### 安装训练依赖 | ||
| ``` | ||
| cd /data/TurboPhysAI | ||
| pip install -r model_examples/Cosmos3/requirements.txt | ||
| pip install lerobot --no-deps | ||
| ``` | ||
|
|
||
| ## 数据与模型权重 | ||
|
|
||
| 按训练任务准备以下资产，路径通过后面的环境变量设置。 | ||
|
|
||
| | 任务 | 数据 | 模型权重 | | ||
| | --- | --- | --- | | ||
| | Generator Nano | BridgeData2 的 `sft_dataset_bridge` | Cosmos3-Nano DCP、Wan2.2 VAE、Cosmos3-Nano processor/tokenizer | | ||
| | Action Policy | DROID LeRobot 数据集 | Cosmos3-Nano DCP、Wan2.2 VAE、Qwen3-VL-8B tokenizer 缓存 | | ||
| | Reasoner | VideoPhy2 或 LLaVA-OV 配方对应数据 | Qwen3-VL 模型目录、Cosmos3-Nano-VLM safetensors | | ||
|
|
||
|
|
||
| ```bash | ||
| cd /data | ||
| wget http://42.228.13.241:18000/docker-images/tencent/Cosmos3/BridgeData2-Subset-Synthetic-Captions.zip | ||
| unzip BridgeData2-Subset-Synthetic-Captions.zip | ||
|
|
||
| wget http://42.228.13.241:18000/docker-images/tencent/Cosmos3/videophysics_official.tar.gz | ||
| tar -xzf videophysics_official.tar.gz -C /data/ | ||
|
|
||
| wget http://42.228.13.241:18000/docker-images/tencent/Cosmos3/Cosmos3-DROID-perf64.tar.gz | ||
| tar -xzf Cosmos3-DROID-perf64.tar.gz -C /data/ | ||
|
|
||
| wget http://42.228.13.241:18000/docker-images/tencent/Cosmos3/Qwen3-VL-8B-Instruct-tokenizer-cache.tar.gz | ||
| tar -xzf Qwen3-VL-8B-Instruct-tokenizer-cache.tar.gz -C /data/ | ||
| ``` | ||
|
|
||
| DROID perf64 子集使用 `/data/Cosmos3-DROID-perf64/keep_ranges_perf64.json` 过滤文件。Tokenizer 缓存解压到 `/data/hf-cache`，不包含 Qwen 模型权重。 | ||
|
|
||
| 将 Cosmos3-Nano HF 权重转换为 DCP 时，在 Cosmos 仓库根目录设置 `WAN_VAE_PATH` 和 `BASE_CHECKPOINT_PATH`，然后执行： | ||
|
|
||
| ```bash | ||
| WAN_VAE_PATH=$WAN_VAE_PATH HF_HUB_OFFLINE=1 \ | ||
| python -m cosmos_framework.scripts.convert_model_to_dcp \ | ||
| -i /path/to/Cosmos3-Nano -o $BASE_CHECKPOINT_PATH | ||
| ``` | ||
|
|
||
| 上游转换脚本会读取 VAE/processor；离线环境需预先准备 Hugging Face 缓存。 | ||
|
|
||
| ## 启动训练 | ||
|
|
||
| 所有命令均在 `cosmos-framework` 根目录执行，训练场景由 `--sft-toml` 选择。 | ||
| 运行环境由 Shell 明确设置，无需 Turbo runtime 配置： | ||
|
|
||
| ```bash | ||
| export PYTHONPATH="$PWD" | ||
| export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 | ||
| export HF_HUB_OFFLINE=1 OMP_NUM_THREADS=1 | ||
| export MIOPEN_FIND_MODE=1 PYTORCH_MIOPEN_SUGGEST_NDHWC=1 | ||
| export TORCH_NCCL_HIGH_PRIORITY=1 | ||
| ``` | ||
|
|
||
| ### Generator Nano | ||
|
|
||
| 设置数据与权重路径后运行： | ||
|
|
||
| ```bash | ||
| export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 | ||
| export DATASET_PATH=/data/BridgeData2-Subset-Synthetic-Captions/sft_dataset_bridge | ||
| export BASE_CHECKPOINT_PATH=path/to/Cosmos3-Nano-DCP | ||
| export WAN_VAE_PATH=path/to/Wan2.2_VAE.pth | ||
| export LOCAL_PROCESSOR_DIR=path/to/Cosmos3-Nano | ||
| export IMAGINAIRE_OUTPUT_ROOT=./outputs/hcu_training | ||
|
|
||
| torchrun --nproc-per-node=8 --master_port=29712 \ | ||
| -m cosmos_framework.scripts.train \ | ||
| --sft-toml=examples/toml/sft_config/vision_sft_nano.toml \ | ||
| -- \ | ||
| job.name=generator_vision_sft_nano_hcu_perf_8card \ | ||
| job.wandb_mode=offline \ | ||
| model.config.vlm_config.tokenizer.pretrained_model_name=$LOCAL_PROCESSOR_DIR \ | ||
| dataloader_train.dataloader.prefetch_factor=16 \ | ||
| trainer.max_iter=70 \ | ||
| checkpoint.save_iter=1000 \ | ||
| trainer.callbacks.compile_tokenizer.enabled=true \ | ||
| trainer.callbacks.compile_tokenizer.compile_after_iterations=3 \ | ||
| 'trainer.callbacks.compile_tokenizer.warmup_resolutions=["256"]' | ||
| ``` | ||
|
|
||
| ### Action Policy（DROID） | ||
|
|
||
| ```bash | ||
| export HF_HOME=/data/cosmos-test/hf-cache | ||
| export BASE_CHECKPOINT_PATH=path/to/Cosmos3-Nano-DCP | ||
| export WAN_VAE_PATH=path/to/Wan2.2_VAE.pth | ||
| export LOCAL_PROCESSOR_DIR=path/to/Cosmos3-Nano | ||
| export DROID_ROOT=/data/Cosmos/datasets/droid_plus_lerobot_640x360_20260412 | ||
| export KEEP_RANGES_PATH=/data/Cosmos/datasets/Cosmos3-DROID-perf64/keep_ranges_perf64.json | ||
| export IMAGINAIRE_OUTPUT_ROOT=./outputs/action_policy | ||
| export TORCHINDUCTOR_WORKER_START=spawn TORCHINDUCTOR_COMPILE_THREADS=4 | ||
| export OMP_NUM_THREADS=4 | ||
|
|
||
| torchrun --nproc_per_node=8 --master_port=29716 \ | ||
| -m cosmos_framework.scripts.train \ | ||
| --sft-toml=examples/toml/sft_config/action_policy_droid_nano.toml -- \ | ||
| job.wandb_mode=disabled job.name=action_policy_sft_nano_hcu_perf_8card_ops \ | ||
| model.config.vlm_config.tokenizer.pretrained_model_name=$LOCAL_PROCESSOR_DIR \ | ||
| model.config.parallelism.data_parallel_shard_degree=8 \ | ||
| model.config.parallelism.data_parallel_replicate_degree=1 \ | ||
| dataloader_train.max_samples_per_batch=8 \ | ||
| dataloader_train.dataloader.batch_size=8 \ | ||
| dataloader_train.dataloader.num_workers=1 \ | ||
| +dataloader_train.dataloader.multiprocessing_context=spawn \ | ||
| dataloader_train.dataloader.prefetch_factor=4 \ | ||
| dataloader_train.dataloader.datasets.droid.dataset.mode=wam \ | ||
| dataloader_train.dataloader.datasets.droid.dataset.video_backend=pyav \ | ||
| dataloader_train.dataloader.datasets.droid.dataset.use_filter_dict=true \ | ||
| dataloader_train.dataloader.datasets.droid.dataset.filter_dict_path="$KEEP_RANGES_PATH" \ | ||
| model.config.rectified_flow_training_config.loss_scale=10.0 \ | ||
| trainer.seed=42 trainer.max_iter=70 \ | ||
| model.config.compile.enabled=true model.config.compile.compile_dynamic=true \ | ||
| trainer.callbacks.compile_tokenizer.enabled=false \ | ||
| checkpoint.save_iter=1000 checkpoint.dcp_async_mode_enabled=false | ||
| ``` | ||
|
|
||
| 本例使用 DROID perf64 子集及对应过滤文件，`MIOPEN_FIND_MODE=1` 由前面的 Shell 设置提供。 | ||
|
|
||
| 一次性性能运行如需关闭 checkpoint 写入，追加 `checkpoint.type._target_=cosmos_framework.checkpoint.load_only.LoadOnlyDistributedCheckpointer`；默认保留正常保存行为。 | ||
|
|
||
| ### Reasoner VideoPhy2 | ||
|
|
||
| ```bash | ||
| export VIDEOPHYSICS_ROOT=/data/videophysics_official | ||
| export VLM_MODEL_NAME=/path/to/Qwen3-VL-8B-Instruct | ||
| export VLM_SAFETENSORS_PATH=/path/to/Cosmos3-Nano-VLM | ||
| export IMAGINAIRE_OUTPUT_ROOT=./outputs/reasoner | ||
|
|
||
| torchrun --nproc_per_node=8 --master_port=29711 \ | ||
| -m cosmos_framework.scripts.train \ | ||
| --sft-toml=examples/toml/sft_config/videophy2_sft_nano.toml -- \ | ||
| job.wandb_mode=disabled job.name=reasoner_videophy2_nano_hcu \ | ||
| model.config.policy.backbone.model_name="$VLM_MODEL_NAME" \ | ||
| model.config.policy.backbone.safetensors_path="$VLM_SAFETENSORS_PATH" \ | ||
| model.config.parallelism.data_parallel_shard_degree=8 \ | ||
| trainer.seed=42 trainer.max_iter=70 scheduler.cycle_lengths=[70] \ | ||
| trainer.callbacks.log_tensor_shape.num_log=-1 checkpoint.save_iter=1000 \ | ||
| model.config.compile.enabled=true model.config.compile.compile_dynamic=true \ | ||
| model.config.activation_checkpointing.mode=none \ | ||
| model.config.parallelism.fsdp_layers_per_group=8 | ||
| ``` | ||
|
|
||
| ### Reasoner LLaVA-OV | ||
|
|
||
| 沿用上面的 VLM 路径。 | ||
|
|
||
| ```bash | ||
| HF_HUB_OFFLINE=0 torchrun --nproc_per_node=8 --master_port=29715 \ | ||
| -m cosmos_framework.scripts.train \ | ||
| --sft-toml=examples/toml/sft_config/llava_ov.toml -- \ | ||
| job.wandb_mode=disabled job.name=reasoner_llava_ov_hcu \ | ||
| model.config.policy.backbone.model_name="$VLM_MODEL_NAME" \ | ||
| model.config.policy.backbone.safetensors_path="$VLM_SAFETENSORS_PATH" \ | ||
| model.config.parallelism.data_parallel_shard_degree=8 data_setting.max_tokens=16000 \ | ||
| trainer.seed=42 trainer.max_iter=70 scheduler.cycle_lengths=[70] \ | ||
| trainer.callbacks.log_tensor_shape.num_log=-1 \ | ||
| checkpoint.save_iter=1000 checkpoint.dcp_async_mode_enabled=false | ||
| ``` | ||
|
|
||
| ## 运行说明 | ||
|
|
||
| - 运行环境以 Shell 中的设置为准，日志与训练产物由 `IMAGINAIRE_OUTPUT_ROOT` 指定。 | ||
| - 修改卡数时，同步调整 `HIP_VISIBLE_DEVICES`、`--nproc_per_node` 和 `data_parallel_shard_degree`。 | ||
| - 调整 Reasoner 训练步数时，同步修改 `trainer.max_iter` 和 `scheduler.cycle_lengths`。 | ||
| - 单项优化通过模型的环境变量或训练配置控制。 | ||
| - 也可直接使用 patch 中的 `examples/launch_sft_*_hcu.sh`，按各脚本说明设置资产与运行参数。 | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.