[WS2][GEMM] Qwen3-8B TP=2 CP=2 deterministic GEMM forward/backward + TP/CP reductions (2 nodes x 2 GPUs, BF16 SM90)

[frank-2077]

WS2 Qwen3-8B Deterministic GEMM Forward/Backward + TP/CP Reductions

Status: Draft

Related work:

• #83: RL-Kernel roadmap
• #108: WS1 numerical contract
• #180: WS1 batch-invariant deterministic GEMM
• WS2 cross-config logprob drift contract

Background

Roadmap #83 places distributed and parallelism invariance under WS2. This issue validates Qwen3-8B GEMM behavior under a fixed TP=2, CP=2 deployment.

Forward and backward are separate communication contracts:

• RowParallel forward produces same-coordinate partial outputs and requires TP SUM.
• ColumnParallel backward produces same-coordinate partial input gradients and requires TP SUM.
• Every TP weight shard is replicated across CP ranks, so its local weight-gradient contributions require CP SUM during backward.

The implementation is built on the fixed-order det_gemm_fwd, det_gemm_da, and det_gemm_db arithmetic introduced by #180. This issue remains GEMM-only: it covers rank-local GEMM arithmetic, fixed tensor ownership, explicit TP/CP reduction operators, and operator-level numerical validation.

The forward track is shared by rollout and training forward. The backward track is training-only and must not be inferred from the forward communication path.

Target Configuration

• Model: Qwen3-8B dense transformer layer
• Architecture: NVIDIA SM90
• Input/output and activation-gradient dtype: BF16
• Local GEMM accumulation dtype: FP32
• Tensor parallelism: TP=2
• Context parallelism: CP=2
• Data parallelism: DP=1
• Physical topology: 2 nodes x 2 GPUs per node, 4 ranks total
• Sharding: even TP and CP shards only
• Execution scope: standalone operator-level forward and backward
• Integration scope: standalone RL-Kernel operators and tests only

Version 1 freezes BF16 local partial outputs and synchronous in-place BF16 reductions. FP32 communication partials and FP32 gradient buffers are separate TODO items.

Qwen3 MLP GEMM Flow

The relevant Transformers implementation is Qwen3MLP.forward:

def forward(self, x):
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
    return down_proj

Expanded into observable GEMM boundaries:

gate   = GEMM(x, W_gate)
up     = GEMM(x, W_up)
hidden = activation(gate) * up
output = GEMM(hidden, W_down)

There are three logical GEMM calls:

1. Gate projection.
2. Up projection.
3. Down projection.

This issue does not pack Gate and Up into one projection. Gate and Up call det_gemm independently. The missing activation/multiply forward and backward boundaries are listed in TODO. Backward tests therefore use independent deterministic dGate, dUp, and dOutput fixtures; this issue does not claim a complete Qwen3 MLP forward or backward result.

Fixed TP/CP Rank Layout

                     TP0                 TP1
CP0 / Node 0         rank0 / GPU0        rank1 / GPU1
CP1 / Node 1         rank2 / GPU0        rank3 / GPU1

Process groups:

TP groups: [rank0, rank1], [rank2, rank3]
CP groups: [rank0, rank2], [rank1, rank3]

CP splits the sequence dimension:

S_local = S / 2
M_local = B * S_local = B * S / 2

This formula assumes DP=1, an even sequence shard, and a dense/padded layout. TP does not split M_local.

Ranks on different nodes process different sequence shards. For the same TP lane, the GEMM weight shard is replicated across CP:

rank0 and rank2 own the TP0 weight shard
rank1 and rank3 own the TP1 weight shard

Ranks on the same node process the same sequence shard and differ because of TP:

• Gate/Up: the GEMM input rows are identical inside a TP group; the weight/output-column shards differ.
• Down: both the hidden-feature shard and the matching weight K-shard differ; the two local outputs are partial sums.

CP does not require activation AllGather inside these GEMM operators. Forward CP AllGather may be used only by the test harness to reconstruct the global sequence.

During backward, the CP group has a different role: it sums local dW contributions for the same replicated TP weight shard. The operator accepts an explicit gradient group; version 1 requires that group to equal the fixed two-rank CP group. It does not infer or create DP/CP groups.

Why Gate/Up Are ColumnParallel and Down Is RowParallel

Use the mathematical Linear layout:

Y[M,N] = X[M,K] @ W[K,N]

ColumnParallel: split W along N/output features
RowParallel:    split W along K/input-reduction features

The selected layout is not the only mathematically possible layout. It is the only layout under the fixed version-1 ownership contract that preserves the TP shard through the nonlinear boundary, returns to the replicated residual stream, and does not add an avoidable TP gather or reduction.

Gate/Up must use aligned ColumnParallel shards

For Qwen3-8B:

X:      [M_local,4096]
W_gate: [4096,12288]
W_up:   [4096,12288]

ColumnParallel splits the intermediate/output dimension:

W_gate = [W_gate_0 | W_gate_1]
W_up   = [W_up_0   | W_up_1]

gate_i = X @ W_gate_i    # [M_local,6144]
up_i   = X @ W_up_i      # [M_local,6144]

Gate and Up must use the same global intermediate-feature offsets. Then every rank can complete the nonlinear boundary locally:

hidden_i = SiLU(gate_i) * up_i

No TP AllReduce or AllGather is required between Gate/Up and Down because SiLU and multiplication are elementwise and do not mix different intermediate-feature coordinates.

The nonlinear boundary is the reason RowParallel Gate/Up is a poor layout. If Gate/Up split the input/reduction dimension, each rank would produce same-coordinate partials:

gate = gate_partial_0 + gate_partial_1
up   = up_partial_0   + up_partial_1

The partials cannot pass through SiLU independently because, in general:

SiLU(a + b) != SiLU(a) + SiLU(b)

Nor can the multiply be distributed over independently incomplete Gate and Up values. Both complete Gate and complete Up must therefore be TP-reduced before activation. That adds two logical reductions over [M_local,12288] tensors, or the same bytes as one packed [M_local,24576] reduction, before Down starts.

ColumnParallel Gate/Up removes those reductions entirely and keeps only [M_local,6144] local shards.

Down must consume the intermediate with RowParallel

After the local nonlinear boundary, each rank owns:

hidden_i: [M_local,6144]

This is already a shard of Down's input/reduction dimension. RowParallel aligns the matching weight shard:

W_down_i: [6144,4096]

output_partial_i = hidden_i @ W_down_i
output = TP_AllReduce_SUM(output_partial_i)

The reduction produces the complete [M_local,4096] residual-stream tensor on both TP ranks. It requires exactly one logical MLP-forward TP reduction.

ColumnParallel Down would require every TP rank to receive the complete [M_local,12288] intermediate before GEMM, so it first needs a TP AllGather. It would then produce an output-hidden shard [M_local,2048], while the fixed residual stream is TP-replicated [M_local,4096]; residual add would require another ownership conversion or a different sequence-parallel contract.

Therefore the fixed forward communication comparison is:

selected layout:
  Gate/Up Column -> local SiLU/multiply -> Down Row
  communication: one TP reduction of [M_local,4096]

RowParallel Gate/Up alternative:
  communication before nonlinearity:
    Gate reduction [M_local,12288]
    Up reduction   [M_local,12288]
  plus a later ownership conversion/Down reduction

ColumnParallel Down alternative:
  gather full intermediate [M_local,12288]
  plus convert output ownership before replicated residual

This issue freezes the selected layout. TP sequence parallel may define a different residual ownership and use AllGather/ReduceScatter, but that is a different operator contract and remains out of scope.

Training TP/CP Communication Guide

The MLP training path must be read in reverse during backward. CP continues to shard token rows; TP continues to shard feature/weight dimensions.

Forward

X_cp: CP-local token rows, TP-replicated hidden [M_local,4096]

Gate/Up ColumnParallel:
  local GEMMs
  no TP collective
  no CP collective

SiLU + multiply:
  local TP feature shards
  no TP collective
  no CP collective

Down RowParallel:
  local output partial
  TP AllReduce over [M_local,4096]
  no CP collective

CP activation AllGather is not part of any MLP GEMM forward. Attention has a separate CP KV exchange, but Attention is outside this issue.

Backward activation gradients

Backward starts with dOutput:[M_local,4096], which is CP-sharded by token and replicated inside the TP group.

Down RowParallel backward:

dHidden_i = dOutput @ W_down_i^T   # [M_local,6144]

Each rank owns a different intermediate-feature shard, so dHidden_i remains local and does not use TP AllReduce.

After the elementwise activation/multiply backward, each rank has aligned dGate_i and dUp_i. Gate/Up ColumnParallel backward produces same-coordinate input-gradient contributions:

dX_gate_partial_i = dGate_i @ W_gate_i^T   # [M_local,4096]
dX_up_partial_i   = dUp_i   @ W_up_i^T     # [M_local,4096]

The communication-minimal complete MLP backward would do:

dX_partial_i = dX_gate_partial_i + dX_up_partial_i
dX = TP_AllReduce_SUM(dX_partial_i)

That is one logical TP reduction over [M_local,4096]. Version 1 of this issue lacks the activation/multiply backward, so PR6 validates the generic ColumnParallel dX communication operator on Gate and Up independently. Combining both local contributions before one TP reduction remains an explicit TODO and must not be claimed by PR6.

CP does not AllGather dX; different CP ranks continue to own gradients for different token rows.

Backward weight gradients

Each CP rank computes weight-gradient contributions from its local token rows:

dW_gate_local_i = X_cp^T      @ dGate_i
dW_up_local_i   = X_cp^T      @ dUp_i
dW_down_local_i = hidden_i^T  @ dOutput

Weights are TP-sharded but replicated across CP. Only ranks with the same TP coordinate own the same weight coordinates, so reductions are:

rank0 + rank2 -> TP0 Gate/Up/Down dW shards
rank1 + rank3 -> TP1 Gate/Up/Down dW shards

These are CP weight-gradient reductions, not CP activation AllGathers and not TP reductions. Under the fixed topology they cross nodes. With DP>1, the same mathematical contributions must be reduced once over the correct combined DP x CP replica group; generic group construction is outside version 1.

Gate, Up, and Down are three logical parameter gradients. PR7 validates them independently. A training framework may pack them into gradient buckets and use fewer physical collective calls, but bucket count is not part of this standalone operator contract.

Training communication summary and PR ownership

Training path	TP communication	CP communication	PR ownership
Gate/Up forward	None	None	PR2/PR4 local validation
Down forward	TP AllReduce of output partial	None	PR2 reference, PR3 operator, PR4 topology
Down backward dHidden	None; remains feature-sharded	None	PR5 reference, PR8 topology
Gate/Up backward dX	TP AllReduce of same-coordinate partials	None	PR5 reference, PR6 operator, PR8 topology
Gate/Up/Down backward dW	None across different TP shards	CP AllReduce across same TP lane	PR5 reference, PR7 operator, PR8 topology
Attention forward/backward	separate Attention TP behavior	KV/attention-gradient exchange	Not in #239

For the complete non-sequence-parallel MLP, the communication-minimal mathematical target is:

forward:
  one TP reduction at Down

backward activation:
  one TP reduction after locally combining Gate and Up dX contributions

backward weights:
  one CP/DP replica reduction per logical weight gradient
  physical calls may be bucketed

Version 1 intentionally tests Gate/Up dX and the three dW tensors independently so every failure has one arithmetic and one process-group explanation. Later fusion/bucketing may reduce call count only after this reference contract passes.

Fixed Forward GEMM Shapes and Communication

det_gemm uses A[M,K] @ B[K,N] -> C[M,N].

GEMM	TP sharding	Global B[K,N]	TP=2 local input	TP=2 local B	Local output	Forward collective
Gate	Column / N	[4096,12288]	[M_local,4096]	[4096,6144]	[M_local,6144]	None
Up	Column / N	[4096,12288]	[M_local,4096]	[4096,6144]	[M_local,6144]	None
Down	Row / K	[12288,4096]	[M_local,6144]	[6144,4096]	[M_local,4096] partial	TP AllReduce SUM

Gate and Up produce different output-column shards and do not sum across TP ranks.

Down computes:

partial_i = det_gemm_fwd(hidden_i, W_down_i)
output    = TP_AllReduce_SUM(partial_i)

The forward communication operator accepts an explicit tp_group. It does not initialize process groups, infer global rank mappings, or use a CP group.

Fixed Backward GEMM Shapes and Communication

Gate/Up ColumnParallel backward

For each Gate or Up projection on TP lane i:

dProjection_i: [M_local,6144]
W_i:           [4096,6144]

dX_partial_i = det_gemm_da(dProjection_i, W_i)
             # [M_local,4096]

dX = TP_AllReduce_SUM(dX_partial_i)

dW_local_i = det_gemm_db(X_local, dProjection_i)
           # [4096,6144]

dW_i = CP_AllReduce_SUM(dW_local_i)

The TP reduction is required because different output-feature shards contribute to the same dX coordinates. The CP reduction is required because CP ranks process different tokens while owning replicas of the same TP weight shard.

Gate and Up are independent in version 1. Combining their rank-local dX contributions before one TP reduction requires the missing activation/multiply backward and is a TODO.

Down RowParallel backward

For TP lane i:

dOutput: [M_local,4096]
W_i:     [6144,4096]

dHidden_i = det_gemm_da(dOutput, W_i)
          # [M_local,6144]
          # remains TP feature-sharded; no TP reduction

dW_local_i = det_gemm_db(hidden_i, dOutput)
           # [6144,4096]

dW_i = CP_AllReduce_SUM(dW_local_i)

RowParallel backward does not TP-reduce dHidden_i: each rank owns a different intermediate-feature shard. Its replicated weight shard still requires CP gradient reduction.

Qwen3 MLP projections have no bias in this scope, so bias-gradient reduction is not included.

Forward and Backward Communication Matrix

Path	Gate/Up ColumnParallel	Down RowParallel
Forward activation	local GEMM; no TP collective	local partial GEMM + TP AllReduce
Backward activation	local dX partial GEMM + TP AllReduce	local TP-sharded dHidden; no TP collective
Backward weight	local dW + CP AllReduce over same TP lane	local dW + CP AllReduce over same TP lane

This matrix is the reason forward and backward must be separate PR tracks. A forward communication wrapper cannot be reused as the backward wrapper by changing tensor names.

Numerical and Topology Boundary

Every communication GEMM must expose or test both sides of its reduction:

forward Down:
  local_output_partial_pre_tp_reduce
  output_post_tp_reduce

backward Gate/Up dX:
  local_dx_partial_pre_tp_reduce
  dx_post_tp_reduce

backward weight gradient:
  local_dw_pre_cp_reduce
  dw_post_cp_reduce

This separates local det_gemm drift from collective dtype, group, stream, and summation-order drift.

In the fixed rank layout:

• TP groups are IntraNode.
• CP groups cross nodes.
• Forward Down and backward Gate/Up dX cover IntraNode TP AllReduce.
• Backward dW covers InterNode CP AllReduce.
• InterNode TP AllReduce remains untested.

Version 1 defines the distributed BF16 oracle explicitly:

local_partial_bf16 = round_bf16(local_fp32_accumulation)
reduced_bf16       = BF16_AllReduce_SUM(local_partial_bf16)

Because local partials are rounded before communication, TP=2/CP=2 output is not required to be bitwise equal to a TP=1/CP=1 GEMM that performs one uninterrupted global reduction. It must match the fixed distributed reference and satisfy the shared #108 tolerance against the global FP32 oracle.

Test reports must state:

COVERED:
  2 nodes x 2 GPUs per node
  TP=2, CP=2, DP=1, even shards
  two IntraNode TP AllReduce groups
  two InterNode CP weight-gradient AllReduce groups

NOT COVERED:
  TP groups crossing nodes
  TP=4, CP=4, or larger world sizes
  DP>1 or a combined DP x CP gradient group
  uneven K/N or sequence shards
  bitwise equality across different NCCL algorithms, protocols, or rank mappings

Goals

1. Reuse a merged or pinned det_gemm_fwd, det_gemm_da, and det_gemm_db implementation from [WS1][kernels] Batch-invariant deterministic GEMM (fwd + bwd) #180; do not create another GEMM arithmetic core.
2. Validate the fixed Qwen3-8B Gate, Up, and Down TP=2 local forward and backward shapes on SM90.
3. Freeze the TP=2, CP=2 rank ownership and independent forward/backward shard references.
4. Implement fixed TP=2 Down forward as local det_gemm_fwd followed by BF16 TP AllReduce SUM.
5. Implement fixed TP=2 Gate/Up backward dX as local det_gemm_da followed by BF16 TP AllReduce SUM.
6. Validate Down backward dHidden remains TP feature-sharded and requires no TP collective.
7. Implement replicated Gate/Up/Down dW as local det_gemm_db followed by BF16 CP AllReduce SUM over the explicit same-TP-lane group.
8. Validate every pre-reduce partial and post-reduce output independently under the shared [WS1] Ground-truth harness + numerical contract for batch-invariant ops #108 numerical contract.
9. Validate the fixed 2-node x 2-GPU deployment and report actual backend, dtype, process groups, NCCL version, rank mapping, and physical topology.
10. Fail closed for unsupported dtype, shape, world size, process group, device, or unavailable deterministic backend.

Non-Goals

• Normalization, activation, elementwise multiply, residual, or complete transformer-layer execution.
• Gate/Up packing, combined Gate/Up backward reduction, or any cross-operator fusion.
• A complete Qwen3 MLP forward or backward result.
• Attention forward/backward or CP KV/attention-gradient communication.
• Arbitrary TP/CP world sizes or a generic shard planner.
• DP>1, generic DP x CP gradient-group construction, DDP/FSDP gradient buckets, or optimizer state/updates.
• InterNode TP AllReduce or TP=4+.
• Uneven TP or CP shards.
• ReduceScatter or TP sequence-parallel ownership.
• FP32 communication partials/gradient buffers, FP16, FP8, quantization, or mixed-quantization paths.
• Autograd integration or framework-owned gradient accumulation.
• vLLM, Megatron, Transformer Engine, FSDP, DeepSpeed, or serving/training integration.
• Performance tuning that weakens or obscures the numerical contract.

Planned PRs

Forward and backward are separate tracks. Each PR adds one new arithmetic or communication boundary and has an independent oracle.

Dependency / #180: stabilize local det_gemm forward and backward primitives

• Land or pin a reviewed commit that exposes det_gemm_fwd, det_gemm_da, and det_gemm_db.
• Use the [WS1] Ground-truth harness + numerical contract for batch-invariant ops #108 BF16 reduction tolerance and preserve the valid bitwise batch/chunk/padding checks.
• Confirm the actual cuda_sm90 backend and fail closed when it is unavailable.
• Cover forward (K,N)=(4096,6144) and (6144,4096) plus the matching local dA and dB shapes.
• Treat dB=A^T@dC as a reduction over local M; do not claim bitwise invariance across different CP partitions or early-rounding boundaries.
• Keep this dependency outside [WS2][GEMM] Qwen3-8B TP=2 CP=2 deterministic GEMM forward/backward + TP/CP reductions (2 nodes x 2 GPUs, BF16 SM90) #239 feature PRs.

Forward track

PR2: fixed TP=2, CP=2 forward shard references

• Freeze the four-rank layout and exact forward shard offsets.
• Validate Gate/Up N-shards by test-only concat against the TP=1 oracle.
• Validate Down K-shards, each local partial, and fixed-order BF16 partial sum without a production collective.
• Use an independent deterministic BF16 Down input fixture because the activation/multiply operator is unavailable.
• Do not initialize NCCL in the arithmetic reference test.

PR3: fixed TP=2 Down forward communication GEMM

• Add a forward operator that accepts hidden_local, weight_local, and an explicit tp_group.
• Compute the BF16 local partial with det_gemm_fwd.
• Run synchronous in-place BF16 AllReduce SUM on the explicit TP group.
• Validate TP world size, local shapes, dtype, device, process-group initialization, and rank ownership.
• First validate one TP group on a single node with two GPUs.
• Compare pre-reduce partials and post-reduce outputs with PR2 references.
• Do not create or use a CP group inside the operator.

PR4: 2-node x 2-GPU forward validation

• Create TP groups [0,1], [2,3] and CP groups [0,2], [1,3] in the test launcher.
• Run Gate and Up local det_gemm_fwd independently on both CP shards.
• Run the same PR3 Down forward operator on both TP groups.
• Use CP AllGather only for optional test reconstruction; verify sequence offsets and shard identity.
• Report repeatability, max/p95/p99 drift, actual backend, CUDA/PyTorch/NCCL versions, group membership, and topology.
• Mark InterNode TP, TP=4+, and uneven shards as NOT COVERED.

Backward track

PR5: fixed TP=2, CP=2 backward shard references

• Freeze independent deterministic dGate, dUp, and dOutput fixtures.
• Validate local Gate/Up dX_partial, Down dHidden_i, and Gate/Up/Down dW_local_i shapes and values without NCCL.
• Build explicit fixed-order TP sums for Gate/Up dX and CP sums for every dW shard.
• Prove Down dHidden_i remains TP feature-sharded and must not be TP-reduced.
• Do not use autograd or create a process group in this reference PR.

PR6: fixed TP=2 ColumnParallel backward dX communication GEMM

• Add a backward-input operator that accepts dOutput_local, weight_local, and an explicit tp_group.
• Compute dX_partial with det_gemm_da.
• Run synchronous in-place BF16 AllReduce SUM on the explicit TP group.
• Validate Gate and Up exact local shapes independently on one node with two GPUs.
• Compare pre-reduce and post-reduce tensors with PR5.
• Do not combine Gate and Up or use a CP group.

PR7: fixed CP=2 replicated-weight dW communication GEMM

• Add a weight-gradient operator that accepts input_local, dOutput_local, and an explicit grad_group.
• Compute dW_local with det_gemm_db.
• Run synchronous in-place BF16 AllReduce SUM over the explicit two-rank group.
• Require the fixed version-1 group to contain CP peers with the same TP coordinate.
• Validate Gate, Up, and Down exact local dW shapes independently.
• Compare local_dw_pre_cp_reduce and dw_post_cp_reduce with PR5.
• Do not infer a global DP x CP group, create optimizer state, or integrate gradient buckets.

PR8: 2-node x 2-GPU backward validation and reproducibility report

• Run PR6 on TP groups [0,1] and [2,3].
• Run PR7 on CP groups [0,2] and [1,3].
• Verify that only same-coordinate TP weight shards are reduced together.
• Verify Down dHidden remains local while Gate/Up dX is TP-reduced.
• Report pre/post-reduce drift, repeatability, actual backend, collective dtype, CUDA/PyTorch/NCCL versions, group membership, rank mapping, and physical topology.
• Mark DP>1, combined DP x CP groups, InterNode TP, TP=4+, and uneven shards as NOT COVERED.
• Do not claim complete Qwen3 MLP backward validation.

TODO

• PreNorm/RMSNorm + GEMM fused variants, still tracked under a GEMM-themed issue.
• Deterministic SwiGLU or SiLU-plus-multiply forward/backward operator for the missing activation boundary.
• Packed Gate/Up projection and combined local Gate/Up dX contribution before one TP reduction.
• Complete composed Qwen3 MLP forward/backward validation after the missing operators exist.
• Attention forward/backward CP communication.
• Generic DP x CP gradient groups, gradient buckets, accumulation, overlap, and sharded optimizers.
• InterNode TP AllReduce and TP=4+/CP=4+ validation.
• FP32 local partial, FP32 AllReduce, and FP32 gradient-buffer modes.
• ReduceScatter and TP sequence-parallel ownership.
• Autograd and framework adapters plus end-to-end train/inference integration.
• Performance optimization after the correctness contract is closed.