[platform]: flux2 & tp#1212
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 Parallelism (TP) support for the Flux2 model, including TP initialization, weight splitting/loading, and rank 0 input/output broadcasting. The review feedback highlights critical issues for hybrid DP+TP and multi-node environments, such as restricting broadcasts to the TP group, using local rank for device ordinals, and correctly identifying rank 0 within the TP group. Additionally, suggestions are made to add safety assertions for attention head divisibility, avoid potential AttributeErrors when clearing device cache, and optimize weight loading performance by broadcasting unsplit weights and slicing locally.
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 _is_rank0(self): | ||
| return not dist.is_initialized() or dist.get_rank() == 0 |
There was a problem hiding this comment.
In a hybrid Data Parallel (DP) + Tensor Parallel (TP) setup, only checking dist.get_rank() == 0 (global rank 0) is insufficient. Each TP group has its own rank 0 (the process responsible for I/O and broadcasting within that TP group). If only global rank 0 loads the text encoder and VAE, other TP groups (e.g., on other DP ranks) will not have their inputs loaded, or they will incorrectly receive global rank 0's inputs if broadcasting across WORLD.
To support DP+TP correctly, _is_rank0 should check if the process is rank 0 within its TP group.
| def _is_rank0(self): | |
| return not dist.is_initialized() or dist.get_rank() == 0 | |
| def _is_rank0(self): | |
| if not dist.is_initialized(): | |
| return True | |
| if self.config.get("tensor_parallel", False): | |
| tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p") | |
| return dist.get_rank(tp_group) == 0 | |
| return dist.get_rank() == 0 |
| def _broadcast_tensor_tree(self, obj, src_obj, src=0): | ||
| if torch.is_tensor(obj): | ||
| if self._is_rank0(): | ||
| obj = src_obj.to(self._rank_device(), non_blocking=True) | ||
| dist.broadcast(obj, src=src) | ||
| return obj | ||
| if isinstance(obj, dict): | ||
| return {key: self._broadcast_tensor_tree(value, src_obj[key] if self._is_rank0() else None, src=src) for key, value in obj.items()} | ||
| if isinstance(obj, list): | ||
| return [self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)] | ||
| if isinstance(obj, tuple): | ||
| return tuple(self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)) | ||
| return obj |
There was a problem hiding this comment.
Update _broadcast_tensor_tree to accept and pass a group argument. This is necessary to restrict the broadcast to the TP group in DP+TP setups, preventing inputs from being incorrectly broadcast across different DP groups.
| def _broadcast_tensor_tree(self, obj, src_obj, src=0): | |
| if torch.is_tensor(obj): | |
| if self._is_rank0(): | |
| obj = src_obj.to(self._rank_device(), non_blocking=True) | |
| dist.broadcast(obj, src=src) | |
| return obj | |
| if isinstance(obj, dict): | |
| return {key: self._broadcast_tensor_tree(value, src_obj[key] if self._is_rank0() else None, src=src) for key, value in obj.items()} | |
| if isinstance(obj, list): | |
| return [self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)] | |
| if isinstance(obj, tuple): | |
| return tuple(self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)) | |
| return obj | |
| def _broadcast_tensor_tree(self, obj, src_obj, src=0, group=None): | |
| if torch.is_tensor(obj): | |
| if self._is_rank0(): | |
| obj = src_obj.to(self._rank_device(), non_blocking=True) | |
| dist.broadcast(obj, src=src, group=group) | |
| return obj | |
| if isinstance(obj, dict): | |
| return {key: self._broadcast_tensor_tree(value, src_obj[key] if self._is_rank0() else None, src=src, group=group) for key, value in obj.items()} | |
| if isinstance(obj, list): | |
| return [self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src, group=group) for idx, value in enumerate(obj)] | |
| if isinstance(obj, tuple): | |
| return tuple(self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src, group=group) for idx, value in enumerate(obj)) | |
| return obj |
| def _broadcast_rank0_payload(self, payload): | ||
| if not self._use_tp_rank0_io(): | ||
| return payload | ||
|
|
||
| meta_list = [self._tensor_meta_tree(payload) if self._is_rank0() else None] | ||
| dist.broadcast_object_list(meta_list, src=0) | ||
| payload_tree = payload if self._is_rank0() else self._materialize_meta_tree(meta_list[0], self._rank_device()) | ||
| payload_tree = self._broadcast_tensor_tree(payload_tree, payload if self._is_rank0() else None, src=0) | ||
| dist.barrier() | ||
| return payload_tree |
There was a problem hiding this comment.
When broadcasting the payload, if group is not specified, dist.broadcast_object_list and dist.broadcast will broadcast across WORLD (all ranks). In a DP+TP setup, this will cause all DP groups to receive the same input from global rank 0, breaking Data Parallelism.
To ensure correct behavior, the broadcast should be restricted to the TP group (tensor_p), and src=0 will correctly refer to the TP rank 0 of each TP group.
| def _broadcast_rank0_payload(self, payload): | |
| if not self._use_tp_rank0_io(): | |
| return payload | |
| meta_list = [self._tensor_meta_tree(payload) if self._is_rank0() else None] | |
| dist.broadcast_object_list(meta_list, src=0) | |
| payload_tree = payload if self._is_rank0() else self._materialize_meta_tree(meta_list[0], self._rank_device()) | |
| payload_tree = self._broadcast_tensor_tree(payload_tree, payload if self._is_rank0() else None, src=0) | |
| dist.barrier() | |
| return payload_tree | |
| def _broadcast_rank0_payload(self, payload): | |
| if not self._use_tp_rank0_io(): | |
| return payload | |
| tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p") if self.config.get("tensor_parallel", False) else None | |
| meta_list = [self._tensor_meta_tree(payload) if self._is_rank0() else None] | |
| dist.broadcast_object_list(meta_list, src=0, group=tp_group) | |
| payload_tree = payload if self._is_rank0() else self._materialize_meta_tree(meta_list[0], self._rank_device()) | |
| payload_tree = self._broadcast_tensor_tree(payload_tree, payload if self._is_rank0() else None, src=0, group=tp_group) | |
| dist.barrier(group=tp_group) | |
| return payload_tree |
| 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() directly as the device ordinal will fail on multi-node setups (e.g., on node 1, global ranks are 8-15, but local GPUs are indexed 0-7, so cuda:8 will raise an invalid device ordinal error).
To support multi-node distributed setups, you should use the local rank (e.g., from LOCAL_RANK environment variable or dist.get_rank() % torch.cuda.device_count()) instead of the global rank.
| 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): | |
| 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) |
| if self.config.get("tensor_parallel", False): | ||
| self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p") | ||
| self.tp_rank = dist.get_rank(self.tp_group) | ||
| self.tp_size = dist.get_world_size(self.tp_group) |
There was a problem hiding this comment.
Add an assertion to ensure that num_attention_heads is divisible by tp_size when tensor parallel is enabled. This prevents potential shape mismatch errors during attention computation.
| if self.config.get("tensor_parallel", False): | |
| self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p") | |
| self.tp_rank = dist.get_rank(self.tp_group) | |
| self.tp_size = dist.get_world_size(self.tp_group) | |
| if self.config.get("tensor_parallel", False): | |
| self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p") | |
| self.tp_rank = dist.get_rank(self.tp_group) | |
| self.tp_size = dist.get_world_size(self.tp_group) | |
| assert self.config.get("num_attention_heads", 24) % self.tp_size == 0, \ | |
| f"num_attention_heads ({self.config.get('num_attention_heads', 24)}) must be divisible by tp_size ({self.tp_size})" |
| for rank_idx in range(self.tp_size): | ||
| if is_weight_loader: | ||
| src_tensor = weight_dict[f"{key}__tp_rank_{rank_idx}"].to(target_device, non_blocking=True) | ||
| else: | ||
| src_tensor = torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device) | ||
| dist.broadcast(src_tensor, src=global_src_rank) | ||
| if rank_idx == self.tp_rank: | ||
| distributed_weight_dict[key].copy_(src_tensor, non_blocking=True) | ||
| del src_tensor |
There was a problem hiding this comment.
Instead of splitting the weights on rank 0 and broadcasting each split separately (which requires self.tp_size broadcasts per parameter and temporary allocations/deallocations on all ranks), consider broadcasting the unsplit weight once to all ranks, and then having each rank perform the local slicing/splitting using _split_weight_for_tp (or _split_bias_for_tp). This reduces the number of collective operations by a factor of self.tp_size, significantly speeding up weight loading and reducing memory fragmentation.
| if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | ||
| del self.text_encoders | ||
| self.text_encoders = None | ||
| torch_device_module.empty_cache() |
There was a problem hiding this comment.
If AI_DEVICE is "cpu", torch_device_module (which is getattr(torch, "cpu", None)) might not have an empty_cache attribute, leading to an AttributeError. It is safer to check if the module has the empty_cache attribute before calling it.
| torch_device_module.empty_cache() | |
| if hasattr(torch_device_module, "empty_cache"): | |
| torch_device_module.empty_cache() |
1 participant
No description provided.