Skip to content

[WS2] TP-invariant reductions (single-GPU vs multi-GPU accumulation order) #109

Description

@Flink-ddd

Part of #83 · refines #102

Make TP=1 and TP>1 produce identical reduction order for the WS1 ops (selected-logprob, masked reductions). The reduction contract and the concrete drift example live in #102. Tolerance reuses the #108 per-dtype threshold table — do not define a separate one.

Relationship to #102 (done vs remaining): #102 is the design/contract + reference issue — it defines the TP-invariant reduction semantics and provides a reference helper on simulated TP shards. Its Non-Goals explicitly exclude real vLLM/Megatron/FSDP integration, and it is only forward/contract-focused. #102 is still open (only the reference scaffold PR #103 is in progress), so the contract is not yet fully frozen. This issue (#109)
owns the parts #102 does not cover: the actual tree-reduce kernel implementation, the TP-wrapper including the backward all-reduce path, real multi-GPU consistency tests (beyond simulated shards), and the performance benchmark. In short: #102 defines the correct answer; #109 builds the working multi-GPU implementation and tests against it. Before starting, confirm with @inaniloquentee which parts of #102's contract are frozen.

Planned PRs:

  • Per-card tree-reduce with a fixed, documented topology (the same shape always takes the same reduction tree); linear reduce allowed as a fast path only when it is bitwise-equivalent (e.g. TP=1).
  • Python TP-wrapper handling post-kernel multi-GPU communication for BOTH forward and backward (backward all-reduce order must also be deterministic).
  • Single-GPU vs multi-GPU data-consistency tests (TP=2 and TP=4, target token on every shard, uneven vocab shard tails).
  • Performance-overhead benchmark (deterministic path vs fast path) + doc.

Depends on: WS1 logprob op working draft. The wrapper/test design can start in W1; full work begins once the WS1 op is available.

Metadata

Metadata

Assignees

Labels

component: distributedTasks involving Ray actor management, cross-node scheduling, and communication synchronization.component: kernelsTasks involving the development of CUDA and Triton underlying operatorsfeaturepriority: highSevere congestion issues require the highest priority for resolution.sprint-0615

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions