[WS2][GEMM][Forward]: implement PR3 TP with FFN and batch-invariance tests - #293
Open
Flink-ddd wants to merge 4 commits into
Open
[WS2][GEMM][Forward]: implement PR3 TP with FFN and batch-invariance tests#293Flink-ddd wants to merge 4 commits into
Flink-ddd wants to merge 4 commits into
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a Qwen3-style tensor-parallel SwiGLU FFN with explicit TP context validation, deterministic local GEMM injection, rank-local weight sharding, autograd-aware collectives, and spawned two-rank Gloo tests. ChangesQwen3 tensor-parallel FFN
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TensorParallelFFN
participant DeterministicGEMM
participant TPGroup
TensorParallelFFN->>TPGroup: copy replicated input for autograd
TensorParallelFFN->>DeterministicGEMM: compute local Gate and Up projections
DeterministicGEMM-->>TensorParallelFFN: return local projections
TensorParallelFFN->>DeterministicGEMM: compute local Down partial
TensorParallelFFN->>TPGroup: all-reduce Down partial
TPGroup-->>TensorParallelFFN: return summed output
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py`:
- Around line 72-85: Add an explicit supported-configuration policy to
FFNContext initialization and the FFN execution boundary: validate the
process-group backend, rank mapping, topology, tensor device, and dtype before
any collective or GEMM. Reject unsupported combinations fail-closed with
diagnostics containing backend, group size/rank, TP size/rank, device, and dtype
metadata. Preserve CPU Gloo as an explicitly supported test configuration, and
ensure validation occurs before execution.
- Around line 249-269: Add validation in TensorParallelFFN construction before
registering gate_weight, up_weight, and down_weight to require all three weights
use the same device and dtype. Reject mismatches consistently for direct
construction, while preserving the existing _make_parameter registration and
shape handling.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 3b19d32b-fb5b-46bc-b35c-d9872323f45d
📒 Files selected for processing (3)
rl_engine/kernels/ops/pytorch/ffn/__init__.pyrl_engine/kernels/ops/pytorch/ffn/tensor_parallel.pytests/test_tensor_parallel_ffn.py
Comment on lines
+72
to
+85
| if not dist.is_available() or not dist.is_initialized(): | ||
| raise RuntimeError( | ||
| "FFNContext(tp_group=...) requires torch.distributed to be initialized." | ||
| ) | ||
| group_size = dist.get_world_size(group=self.tp_group) | ||
| group_rank = dist.get_rank(group=self.tp_group) | ||
| size = group_size if self.tp_size is None else int(self.tp_size) | ||
| rank = group_rank if self.tp_rank is None else int(self.tp_rank) | ||
| if size != group_size: | ||
| raise ValueError( | ||
| f"ctx.tp_size={size} does not match tp_group world size={group_size}." | ||
| ) | ||
| if rank != group_rank: | ||
| raise ValueError(f"ctx.tp_rank={rank} does not match tp_group rank={group_rank}.") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Validate supported TP execution configurations before execution.
FFNContext accepts every initialized process group. It does not validate or report the backend, rank mapping, topology, or supported device and dtype combination. A production BF16 execution can therefore enter an unsupported configuration and fail later in GEMM or all_reduce without the required fail-closed diagnostic.
Add an explicit supported-configuration policy. Validate the process-group backend in FFNContext. Validate tensor device and dtype at the FFN execution boundary. Report the required configuration metadata. Keep CPU Gloo support explicit for the test path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py` around lines 72 - 85,
Add an explicit supported-configuration policy to FFNContext initialization and
the FFN execution boundary: validate the process-group backend, rank mapping,
topology, tensor device, and dtype before any collective or GEMM. Reject
unsupported combinations fail-closed with diagnostics containing backend, group
size/rank, TP size/rank, device, and dtype metadata. Preserve CPU Gloo as an
explicitly supported test configuration, and ensure validation occurs before
execution.
Comment on lines
+249
to
+269
| self.gate_weight = self._make_parameter( | ||
| gate_weight, | ||
| (self.local_intermediate_size, hidden_size), | ||
| "gate_weight", | ||
| device=device, | ||
| dtype=dtype, | ||
| ) | ||
| self.up_weight = self._make_parameter( | ||
| up_weight, | ||
| (self.local_intermediate_size, hidden_size), | ||
| "up_weight", | ||
| device=device, | ||
| dtype=dtype, | ||
| ) | ||
| self.down_weight = self._make_parameter( | ||
| down_weight, | ||
| (hidden_size, self.local_intermediate_size), | ||
| "down_weight", | ||
| device=device, | ||
| dtype=dtype, | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject mixed projection weight devices and dtypes.
shard_qwen3_ffn_weights rejects inconsistent full weights, but direct TensorParallelFFN(...) construction bypasses that validation. Mixed Gate, Up, and Down weights create a module that fails later in the injected GEMM, NativeSwiGLUOp, or Down projection.
Validate that all three registered parameters share one device and one dtype.
Proposed fix
self.down_weight = self._make_parameter(
down_weight,
(hidden_size, self.local_intermediate_size),
"down_weight",
device=device,
dtype=dtype,
)
+ if not (
+ self.gate_weight.device == self.up_weight.device == self.down_weight.device
+ ):
+ raise ValueError("gate_weight, up_weight, and down_weight must share one device.")
+ if not (
+ self.gate_weight.dtype == self.up_weight.dtype == self.down_weight.dtype
+ ):
+ raise ValueError("gate_weight, up_weight, and down_weight must share one dtype.")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| self.gate_weight = self._make_parameter( | |
| gate_weight, | |
| (self.local_intermediate_size, hidden_size), | |
| "gate_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| self.up_weight = self._make_parameter( | |
| up_weight, | |
| (self.local_intermediate_size, hidden_size), | |
| "up_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| self.down_weight = self._make_parameter( | |
| down_weight, | |
| (hidden_size, self.local_intermediate_size), | |
| "down_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| self.gate_weight = self._make_parameter( | |
| gate_weight, | |
| (self.local_intermediate_size, hidden_size), | |
| "gate_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| self.up_weight = self._make_parameter( | |
| up_weight, | |
| (self.local_intermediate_size, hidden_size), | |
| "up_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| self.down_weight = self._make_parameter( | |
| down_weight, | |
| (hidden_size, self.local_intermediate_size), | |
| "down_weight", | |
| device=device, | |
| dtype=dtype, | |
| ) | |
| if not ( | |
| self.gate_weight.device == self.up_weight.device == self.down_weight.device | |
| ): | |
| raise ValueError("gate_weight, up_weight, and down_weight must share one device.") | |
| if not ( | |
| self.gate_weight.dtype == self.up_weight.dtype == self.down_weight.dtype | |
| ): | |
| raise ValueError("gate_weight, up_weight, and down_weight must share one dtype.") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/ops/pytorch/ffn/tensor_parallel.py` around lines 249 - 269,
Add validation in TensorParallelFFN construction before registering gate_weight,
up_weight, and down_weight to require all three weights use the same device and
dtype. Reject mismatches consistently for direct construction, while preserving
the existing _make_parameter registration and shape handling.
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
resolved #239 (PR3 Forward track)
Overview
This PR implements PR3: Tensor Parallel (TP) with FFN for the Qwen3-8B architecture, establishing the foundational orchestration, weight slicing, autograd communication mapping, and batch-invariance verification framework.
Key Changes
tensor_parallel.py:
Introduced FFNContext and implemented TP weight slicing for gate_weight, up_weight, and down_weight.
Implemented TensorParallelFFN orchestration module with customized autograd communication mapping.
test_tensor_parallel_ffn.py: Added TP=2 Gloo multi-process correctness, backward communication placement, and batch/padding invariance tests.
all_reduce Placement & Logic:
Constraints & Guardrails:
Verification & Test Results
Local Validation
tests/test_tensor_parallel_ffn.py: 5 passed
tests/test_swiglu.py: 22 passed (81 skipped)
ruff linting: Passed
GPU Hardware Validation (Pending)
Reserved section for multi-GPU hardware test metrics and execution logs:
Summary by CodeRabbit
New Features
Tests