megabatch: harden global_comm_dim_size contract#77
Merged
Conversation
Follow-up to #75. Three small hardening fixes for the new ``global_comm_dim_size`` parameter: 1. Replace ``assert`` with explicit raise. Both the missing-arg check and the FSDP2 sharding-invariant check disappear under ``python -O``; without them, a missing arg degenerates into a confusing ``TypeError`` deep inside the ceil-div, and a violated invariant silently produces the wrong pad size. Use ``ValueError`` / ``RuntimeError`` so the contract holds regardless of optimization. 2. Drop the ``= None`` default on ``global_comm_dim_size``. The signature said optional but the function actually requires it on the sharded path; all three production callers (``dion2``, ``muon``, ``normuon``) already pass it. Forcing explicit pass-through prevents silent regressions where a new caller forgets the arg. 3. Add ``isinstance(X[0], DTensor)`` checks at the three caller extraction sites. The comment "X[0] is still a DTensor here" is a load-bearing precondition (``.shape`` returns the global size only for DTensor); a future refactor that moves ``to_local(X)`` above the extraction would silently produce a local-size-based pad target. The ``RuntimeError`` in megabatch_base would eventually catch it, but from a confusing site -- the TypeError at the caller fails fast with a clear message. Also switch the ceil-div from ``-(-x // y)`` to the more readable ``(x + y - 1) // y``, and tighten the NOTE comment that referenced the removed assert. Tests: add a CPU-only single-rank gloo test pinning the new ValueError contract, so ``python -O`` regressions are caught at the unit-test level.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #75 addressing three review nits without changing the perf win or the API contract.
Changes
assert→ explicitraiseinmegabatch_orthogonalize_async. Both the missing-arg check and the FSDP2 sharding-invariant check vanish underpython -O; without them, a missing arg degenerates into a confusingTypeErrordeep inside the ceil-div, and a violated invariant silently produces the wrong pad size. NowValueError/RuntimeErrorso the contract holds regardless of optimization level.Drop the
= Nonedefault onglobal_comm_dim_size. The signature said optional but the function requires it on the sharded path; all three production callers (dion2,muon,normuon) already pass it. Forcing explicit pass-through prevents silent regressions if a new caller forgets the arg.isinstance(X[0], DTensor)checks at the three caller extraction sites. The comment "X[0] is still a DTensor here" is a load-bearing precondition —.shapereturns the global size only for DTensor. A future refactor that movesto_local(X)above the extraction would silently produce a local-size-based pad target. TheRuntimeErrorinmegabatch_basewould eventually catch it, but from a confusing site; the newTypeErrorat the caller fails fast with a clear message.Also: switch the ceil-div from
-(-x // y)to(x + y - 1) // yfor readability, and update the NOTE comment that referenced the removed assert.Tests
Adds a CPU-only single-rank gloo test that drives the sharded branch with
global_comm_dim_size=Noneand pins the newValueError. Catchespython -O/ future-refactor regressions at the unit level without burning a GPU slot.Existing test_megabatch_empty_shard tests remain unchanged.