Conversation
![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 tensor-parallel (TP) support for Flux2 and Wan2.2 models, including TP configuration files, weight/bias splitting utilities, and a TPRunnerMixin to handle rank-0 I/O broadcasting. The review feedback highlights several critical issues: using dist.get_rank() for device ordinals assumes a single-node setup and should be replaced with LOCAL_RANK to support multi-node clusters; specifying group=self.tp_group in dist.broadcast during weight loading will cause crashes when DP > 1 and should be broadcast globally; and the distributed inference script for Flux2 incorrectly references a non-existent configuration file instead of the newly added flux2_dev_tp.json.
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.
| def _rank_device(self): | ||
| ai_device = global_var.AI_DEVICE | ||
| if ai_device is None: | ||
| return torch.device("cpu") | ||
| if dist.is_initialized(): | ||
| return torch.device(f"{ai_device}:{dist.get_rank()}") | ||
| return torch.device(ai_device) |
There was a problem hiding this comment.
Using dist.get_rank() as the device ordinal assumes a single-node multi-GPU setup. In multi-node distributed environments, dist.get_rank() returns the global rank (which can exceed the number of GPUs on a single node), causing invalid device ordinal runtime errors. Using the LOCAL_RANK environment variable ensures correct local GPU mapping on multi-node clusters.
| def _rank_device(self): | |
| ai_device = global_var.AI_DEVICE | |
| if ai_device is None: | |
| return torch.device("cpu") | |
| if dist.is_initialized(): | |
| return torch.device(f"{ai_device}:{dist.get_rank()}") | |
| return torch.device(ai_device) | |
| def _rank_device(self): | |
| import os | |
| ai_device = global_var.AI_DEVICE | |
| if ai_device is None: | |
| return torch.device("cpu") | |
| if dist.is_initialized(): | |
| local_rank = int(os.environ.get("LOCAL_RANK", 0)) | |
| return torch.device(f"{ai_device}:{local_rank}") | |
| return torch.device(ai_device) |
| def _rank_device(self): | ||
| if dist.is_initialized(): | ||
| return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") | ||
| return torch.device(AI_DEVICE) |
There was a problem hiding this comment.
Using dist.get_rank() as the device ordinal assumes a single-node multi-GPU setup. In multi-node distributed environments, dist.get_rank() returns the global rank (which can exceed the number of GPUs on a single node), causing invalid device ordinal runtime errors. Using the LOCAL_RANK environment variable ensures correct local GPU mapping on multi-node clusters.
| def _rank_device(self): | |
| if dist.is_initialized(): | |
| return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") | |
| return torch.device(AI_DEVICE) | |
| def _rank_device(self): | |
| import os | |
| if dist.is_initialized(): | |
| local_rank = int(os.environ.get("LOCAL_RANK", 0)) | |
| return torch.device(f"{AI_DEVICE}:{local_rank}") | |
| return torch.device(AI_DEVICE) |
| def _rank_device(self): | ||
| if dist.is_initialized(): | ||
| return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") | ||
| return torch.device(AI_DEVICE) |
There was a problem hiding this comment.
Using dist.get_rank() as the device ordinal assumes a single-node multi-GPU setup. In multi-node distributed environments, dist.get_rank() returns the global rank (which can exceed the number of GPUs on a single node), causing invalid device ordinal runtime errors. Using the LOCAL_RANK environment variable ensures correct local GPU mapping on multi-node clusters.
| def _rank_device(self): | |
| if dist.is_initialized(): | |
| return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") | |
| return torch.device(AI_DEVICE) | |
| def _rank_device(self): | |
| import os | |
| if dist.is_initialized(): | |
| local_rank = int(os.environ.get("LOCAL_RANK", 0)) | |
| return torch.device(f"{AI_DEVICE}:{local_rank}") | |
| return torch.device(AI_DEVICE) |
| for key in sorted(synced_meta.keys()): | ||
| m = synced_meta[key] | ||
| if m["is_tp"]: | ||
| for r in range(self.tp_size): | ||
| buf = processed[f"{key}__tp_{r}"].to(target_device) if is_weight_loader else torch.empty(m["shape"], dtype=m["dtype"], device=target_device) | ||
| dist.broadcast(buf, src=src_rank, group=self.tp_group) | ||
| if r == self.tp_rank: | ||
| distributed[key].copy_(buf) | ||
| del buf | ||
| else: | ||
| if is_weight_loader: | ||
| distributed[key].copy_(processed[key].to(target_device)) | ||
| dist.broadcast(distributed[key], src=src_rank, group=self.tp_group) | ||
|
|
||
| return distributed |
There was a problem hiding this comment.
Specifying group=self.tp_group in dist.broadcast will cause a crash or hang when DP > 1 (multiple TP groups) because global rank 0 (the source rank) is not a member of other TP groups. Since only global rank 0 loads the weights, the broadcast must happen across the default global group (world), matching dist.broadcast_object_list and flux2/model.py.
| for key in sorted(synced_meta.keys()): | |
| m = synced_meta[key] | |
| if m["is_tp"]: | |
| for r in range(self.tp_size): | |
| buf = processed[f"{key}__tp_{r}"].to(target_device) if is_weight_loader else torch.empty(m["shape"], dtype=m["dtype"], device=target_device) | |
| dist.broadcast(buf, src=src_rank, group=self.tp_group) | |
| if r == self.tp_rank: | |
| distributed[key].copy_(buf) | |
| del buf | |
| else: | |
| if is_weight_loader: | |
| distributed[key].copy_(processed[key].to(target_device)) | |
| dist.broadcast(distributed[key], src=src_rank, group=self.tp_group) | |
| return distributed | |
| for key in sorted(synced_meta.keys()): | |
| m = synced_meta[key] | |
| if m["is_tp"]: | |
| for r in range(self.tp_size): | |
| buf = processed[f"{key}__tp_{r}"].to(target_device) if is_weight_loader else torch.empty(m["shape"], dtype=m["dtype"], device=target_device) | |
| dist.broadcast(buf, src=src_rank) | |
| if r == self.tp_rank: | |
| distributed[key].copy_(buf) | |
| del buf | |
| else: | |
| if is_weight_loader: | |
| distributed[key].copy_(processed[key].to(target_device)) | |
| dist.broadcast(distributed[key], src=src_rank) | |
| return distributed |
| --model_path $model_path \ | ||
| --prompt "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom." \ | ||
| --save_result_path "${lightx2v_path}/save_results/flux2_dev.png" \ | ||
| --config_json "${lightx2v_path}/configs/flux2/flux2_dev_dist.json" |
There was a problem hiding this comment.
The script references flux2_dev_dist.json which is not added in this PR, whereas flux2_dev_tp.json is added and contains the matching "tensor_p_size": 4 configuration. Update the script to use flux2_dev_tp.json.
| --config_json "${lightx2v_path}/configs/flux2/flux2_dev_dist.json" | |
| --config_json "${lightx2v_path}/configs/flux2/flux2_dev_tp.json" |
1 participant
No description provided.