Skip to content

[DistMuon] Drop the foreach-copy compatibility fallback - #4271

Draft
weifengpy wants to merge 2 commits into
gh/weifengpy/65/basefrom
gh/weifengpy/65/head
Draft

[DistMuon] Drop the foreach-copy compatibility fallback#4271
weifengpy wants to merge 2 commits into
gh/weifengpy/65/basefrom
gh/weifengpy/65/head

Conversation

@weifengpy

@weifengpy weifengpy commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

_foreach_copy_or_fallback_ checks four conditions before using
torch._foreach_copy_ and copies pair-by-pair when any fails. Three of the
four are unnecessary: the op casts between dtypes, crosses devices, and handles
non-contiguous views exactly as Tensor.copy_ does. Only unequal shapes are
a real difference -- copy_ broadcasts, the foreach op raises -- and callers
pair each destination with a source built from the same _TensorRegion, so
their shapes always agree. When they do not, the op's own error is clearer than
silently taking a slower path.

Measured on Moonlight-16B-A3B, DP-shard 8 with EP 4 over 8xH100, 10 steps: the
fallback ran 0 times out of 3220 multi-pair calls, and no condition failed
once. Every batched call held exactly 8 uniform pairs.

Rename to _copy_region_views_, since there is no longer a fallback to
name.

Test plan:

  • pytest tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py -- 6
    passed. The dtype test now asserts the foreach path is taken and still casts,
    which is the evidence for removing the guard; a new case covers the shape
    mismatch the op rejects.
  • Bit-identical numerics against the parent commit on the config above with
    --debug.seed 42 --debug.deterministic: step 1 loss 12.47077 grad_norm 1.1954,
    step 10 loss 10.69065 grad_norm 1.9209.

[ghstack-poisoned]
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 21, 2026
[ghstack-poisoned]
weifengpy added a commit that referenced this pull request Aug 21, 2026
``_foreach_copy_or_fallback_`` checks four conditions before using
``torch._foreach_copy_`` and copies pair-by-pair when any fails. Three of the
four are unnecessary: the op casts between dtypes, crosses devices, and handles
non-contiguous views exactly as ``Tensor.copy_`` does. Only unequal shapes are
a real difference -- ``copy_`` broadcasts, the foreach op raises -- and callers
pair each destination with a source built from the same ``_TensorRegion``, so
their shapes always agree. When they do not, the op's own error is clearer than
silently taking a slower path.

Measured on Moonlight-16B-A3B, DP-shard 8 with EP 4 over 8xH100, 10 steps: the
fallback ran 0 times out of 3220 multi-pair calls, and no condition failed
once. Every batched call held exactly 8 uniform pairs.

Rename to ``_batched_copy_``: there is no longer a fallback to name, and
reusing ``_foreach_copy_`` would shadow the torch op this calls and the tests
patch.

Test plan:
- pytest tests/unit_tests/flex_shard/test_optimizer_reshard_runtime.py -- 6
  passed. The dtype test now asserts the foreach path is taken and still casts,
  which is the evidence for removing the guard; a new case covers the shape
  mismatch the op rejects.
- Bit-identical numerics against the parent commit on the config above with
  --debug.seed 42 --debug.deterministic: step 1 loss 12.47077 grad_norm 1.1954,
  step 10 loss 10.69065 grad_norm 1.9209.

ghstack-source-id: 3379b3e
Pull-Request: #4271
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 Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant