Skip to content

trainer: ddp_bf16_compress, DDP with a bf16 gradient all-reduce - #667

Draft
RyanPCo wants to merge 1 commit into
ryanco/eval-video-fpsfrom
ryanco/ddp-bf16-hook
Draft

RyanPCo wants to merge 1 commit into
ryanco/eval-video-fpsfrom
ryanco/ddp-bf16-hook

Conversation

@RyanPCo

@RyanPCo RyanPCo commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor

On the lab's 8 x A40 nodes -- no NVLink, two NUMA islands joined by
cross-socket links -- the per-step gradient all-reduce, not the math, sets the
step time once a model is large enough. Measured 2026-09-16 on a ~1.2 B
-parameter policy: 6.4 s/step at 8 x batch 16 against 1.4 s/step on a single
GPU, and 3.6 s/step with the reduction in bf16. Only the reduction is bf16;
master weights and the optimizer stay fp32.

It is plain DDP otherwise -- no unused-parameter detection -- so a model that
leaves a trainable tensor without a gradient needs ddp_pi's
ddp_find_unused_parameters_true instead and will crash under this one. The
config says so, because nothing else would stop you selecting it for Pi.

The hook is a zero-argument factory rather than a Hydra _partial_ because
Lightning logs the hook's qualname, which a functools.partial does not
carry. make_bf16_compress_hook is named for what it is: torch's own hook is
bf16_compress_hook, and sharing that name made the factory look like the hook.

test_eval_determinism's trainer parametrization now enumerates the config group
instead of a hardcoded list, so a new trainer config cannot land without the
gradient-clip knob being checked.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo commented Sep 19, 2026 •

Copy link
Copy Markdown
Contributor Author

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

This stack of pull requests is managed by Graphite. Learn more about stacking.

@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 4bbefe4 to f0507d5 Compare September 20, 2026 19:09
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch 3 times, most recently from c5fd130 to d54b1cf Compare September 22, 2026 13:18
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from b6f9d4d to 26eb53e Compare October 12, 2026 21:50
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch from 3adf25b to 57330fa Compare October 12, 2026 23:19
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 26eb53e to 362dd61 Compare October 12, 2026 23:19
@RyanPCo RyanPCo mentioned this pull request Oct 19, 2026
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 362dd61 to 65fea4d Compare October 19, 2026 19:44
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch from 57330fa to 3b31bf1 Compare October 19, 2026 19:44
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 65fea4d to 9d8c5a2 Compare October 19, 2026 20:22
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch from 3b31bf1 to 383f798 Compare October 19, 2026 20:22
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch from 383f798 to 31aba89 Compare October 30, 2026 18:27
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 9d8c5a2 to 2679b17 Compare October 30, 2026 18:27

@RyanPCo RyanPCo left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 nit.

# gradient reduction, not the math, sets the step time. On the lab's 8 x A40
# nodes -- no NVLink, two NUMA islands joined by cross-socket links -- a ~1.2 B
# -parameter policy measured 6.4 s/step at 8 x batch 16 against 1.4 s/step on a
# single GPU, and 3.6 s/step once the reduction was halved (2026-09-16). Only

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit (leftover-experiment): The dated 8xA40 benchmark (6.4 s/step vs 3.6 s/step, 2026-09-16) belongs in the PR description, not the config. 'bf16 gradient all-reduce for comm-bound DDP; master weights and optimizer stay fp32' is enough. On line 15, a one-line note that the parent's strategy: ddp string has to become a DDPStrategy object to accept ddp_comm_hook would explain the re-specified target.

On the lab's 8 x A40 nodes -- no NVLink, two NUMA islands joined by
cross-socket links -- the per-step gradient all-reduce, not the math, sets the
step time once a model is large enough. Measured 2026-09-16 on a ~1.2 B
-parameter policy: 6.4 s/step at 8 x batch 16 against 1.4 s/step on a single
GPU, and 3.6 s/step with the reduction in bf16. Only the reduction is bf16;
master weights and the optimizer stay fp32.

It is plain DDP otherwise -- no unused-parameter detection -- so a model that
leaves a trainable tensor without a gradient needs ddp_pi's
`ddp_find_unused_parameters_true` instead and will crash under this one. The
config says so, because nothing else would stop you selecting it for Pi.

The hook is a zero-argument factory rather than a Hydra `_partial_` because
Lightning logs the hook's __qualname__, which a functools.partial does not
carry. `make_bf16_compress_hook` is named for what it is: torch's own hook is
`bf16_compress_hook`, and sharing that name made the factory look like the hook.

test_eval_determinism's trainer parametrization now enumerates the config group
instead of a hardcoded list, so a new trainer config cannot land without the
gradient-clip knob being checked.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01THTuWdgk6M6pZ5Ya7ihZEF
@RyanPCo
RyanPCo force-pushed the ryanco/eval-video-fps branch from 2679b17 to bbd5c3e Compare October 30, 2026 23:05
@RyanPCo
RyanPCo force-pushed the ryanco/ddp-bf16-hook branch from 31aba89 to bb5001f Compare October 30, 2026 23:05
@RyanPCo

RyanPCo commented Oct 30, 2026

Copy link
Copy Markdown
Contributor Author

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

  • Dated cluster benchmark in config comment: Removed the 8xA40 numbers. Added a one-line comment explaining why the strategy is re-specified as a DDPStrategy object.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant