feat(phyai): Support Cosmos3, TP and CFG Parallel#37
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (26)
📝 WalkthroughWalkthroughThis PR adds a multi-GPU ("wn") tensor-parallel and CFG-parallel execution path for Cosmos3 models, including a new ChangesCosmos3 WN Multi-GPU Path
Estimated code review effort: 4 (Complex) | ~75 minutes Repository Housekeeping
Estimated code review effort: 1 (Trivial) | ~5 minutes Sequence Diagram(s)sequenceDiagram
participant Script as run_cosmos3_wn.py
participant Engine
participant Scheduler as Cosmos3T2VWNScheduler
participant VAE as Cosmos3WanVAE
participant Disk
Script->>Engine: construct(cosmos3_wn, ParallelConfig cfg_size/tp_size)
Script->>Engine: step(request) on all ranks
Engine->>Scheduler: run UniPC denoise loop
alt cfg_size > 1
Scheduler->>Scheduler: cond/uncond branches on cfg axis
Scheduler->>Scheduler: all_gather velocities across cfg
else
Scheduler->>Scheduler: sequential cond/uncond on single rank
end
Scheduler->>VAE: decode_parallel(latents) if tp/cfg > 1
VAE-->>Scheduler: blended pixel tiles
Scheduler-->>Engine: video/sound result
Engine-->>Script: identical result across ranks
Script->>Disk: save mp4 (rank 0 only)
Possibly related PRs
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces multi-GPU tensor-parallel and classifier-free-guidance (CFG) parallel support for Cosmos3 generation and policy models, including new orchestrators, parallel VAE decoding, and fused QKV projections. Key feedback includes adding guards to prevent potential TypeError and incorrect slicing when optional parameters are None in the schedulers, and using the global RANK instead of LOCAL_RANK in the example scripts to ensure correct behavior in multi-node distributed environments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| .standard_normal((batch, chunk, ad)) | ||
| .astype("float32") | ||
| ).to(dev, dt) | ||
| action[:, :, raw:] = 0.0 |
There was a problem hiding this comment.
If raw (i.e., request.raw_action_dim) is None, the slice raw: becomes None:, which is equivalent to 0: in Python slicing. This will silently zero out the entire action tensor instead of just the padding. A guard should be added to ensure raw is not None before zeroing out the padding. This also applies to other lines where raw: is used for slicing (lines 177, 242, 249, 283, 290).
| action[:, :, raw:] = 0.0 | |
| if raw is not None: | |
| action[:, :, raw:] = 0.0 |
| .contiguous() | ||
| ) | ||
|
|
||
| cond_idx = list(request.cond_frame_indexes) |
There was a problem hiding this comment.
If request.cond_frame_indexes is None, calling list(request.cond_frame_indexes) will raise a TypeError: 'NoneType' object is not iterable. We should add a guard to handle the None case gracefully.
| cond_idx = list(request.cond_frame_indexes) | |
| cond_idx = list(request.cond_frame_indexes) if request.cond_frame_indexes is not None else [] |
| f"launched under torchrun (WORLD_SIZE={env_world}) but --cfg*--tp is 1; " | ||
| f"set --cfg/--tp to use all ranks." | ||
| ) | ||
| return local_rank, cfg, tp, local_rank == 0 |
There was a problem hiding this comment.
In a multi-node distributed environment, LOCAL_RANK will be 0 on the first GPU of every node. Consequently, is_main (which is local_rank == 0) will evaluate to True on multiple nodes, causing multiple processes to write to the same output file or print duplicate logs. It is safer and more standard to use the global RANK environment variable to determine the main process.
| return local_rank, cfg, tp, local_rank == 0 | |
| global_rank = int(os.environ.get("RANK", "0")) | |
| return local_rank, cfg, tp, global_rank == 0 |
| f"launched under torchrun (WORLD_SIZE={env_world}) but --cfg*--tp is 1; " | ||
| f"set --cfg/--tp to use all ranks." | ||
| ) | ||
| return local_rank, cfg, tp, local_rank == 0 |
There was a problem hiding this comment.
In a multi-node distributed environment, LOCAL_RANK will be 0 on the first GPU of every node. Consequently, is_main (which is local_rank == 0) will evaluate to True on multiple nodes, causing multiple processes to write to the same output file or print duplicate logs. It is safer and more standard to use the global RANK environment variable to determine the main process.
| return local_rank, cfg, tp, local_rank == 0 | |
| global_rank = int(os.environ.get("RANK", "0")) | |
| return local_rank, cfg, tp, global_rank == 0 |
1 participant
Summary by CodeRabbit
New Features
Documentation
Bug Fixes
Chores