Skip to content

Make OL gradient A2A pipeline-triggered (#4498) - #4498

Closed
spmex wants to merge 1 commit into
meta-pytorch:mainfrom
spmex:export-D110148587
Closed

Make OL gradient A2A pipeline-triggered (#4498)#4498
spmex wants to merge 1 commit into
meta-pytorch:mainfrom
spmex:export-D110148587

Conversation

@spmex

@spmex spmex commented Aug 6, 2026

Copy link
Copy Markdown

Summary:

Context

PECAll2AllSeqWait.backward eagerly fired the OL gradient AllToAll, so gradient communication was not fully under pipeline control -- which the compute-ahead pipeline needs for overlap and for the deferred-NOL schedule.

Approach

  • Remove the eager OL dist() from PECAll2AllSeqWait.backward so backward only splits the gradient and creates both OL/NOL PECGradientApply instances. The pipeline now starts both AllToAlls via grad_dist, making gradient communication fully pipeline-controlled (OL before optimizer, NOL deferred to the next batch).
  • No test change needed: the existing _pec_backward already calls grad_dist(is_overlapped=True) before waiting OL, so the OL path still triggers dist().

Differential Revision: D110148587

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 6, 2026
@meta-codesync

meta-codesync Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

@spmex has exported this pull request. If you are a Meta employee, you can view the originating Diff in D110148587.

Summary:

## Context

`PECAll2AllSeqWait.backward` eagerly fired the OL gradient AllToAll, so gradient communication was not fully under pipeline control -- which the compute-ahead pipeline needs for overlap and for the deferred-NOL schedule.

## Approach

- Remove the eager OL `dist()` from `PECAll2AllSeqWait.backward` so backward only splits the gradient and creates both OL/NOL `PECGradientApply` instances. The pipeline now starts both AllToAlls via `grad_dist`, making gradient communication fully pipeline-controlled (OL before optimizer, NOL deferred to the next batch).
- No test change needed: the existing `_pec_backward` already calls `grad_dist(is_overlapped=True)` before waiting OL, so the OL path still triggers `dist()`.

___

Differential Revision: D110148587
@meta-codesync meta-codesync Bot changed the title Make OL gradient A2A pipeline-triggered Make OL gradient A2A pipeline-triggered (#4498) Aug 6, 2026
@spmex
spmex force-pushed the export-D110148587 branch from fd8af09 to 66bb5aa Compare August 6, 2026 17:00
@meta-codesync meta-codesync Bot closed this in 9c7f6b1 Aug 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant