[PyTorch] Add newton_schulz_tp optimizer step function - #2920
Conversation
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThe PR adds tensor-parallel Newton–Schulz orthogonalization and restores compatibility for the former module path.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["newton_schulz_tp(x, partition_dim, tp_mode)"] --> B{"Replicated input?"}
B -- Yes --> C["Shard larger dimension"]
B -- No --> D{"Duplicated mode?"}
D -- Yes --> E["Gather full tensor"]
D -- No --> F{"Row partition?"}
F -- Yes --> G["Transpose local shard"]
F -- No --> H["Use column shard directly"]
C --> I["Distributed Newton–Schulz kernel"]
E --> I
G --> I
H --> I
I --> J["Gather/transpose/slice as needed"]
J --> K["Copy result into x"]
Reviews (9): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
| def step(self, closure=None): | ||
| """Perform a single optimization step.""" | ||
| loss = None | ||
| if closure is not None: | ||
| loss = closure() | ||
|
|
There was a problem hiding this comment.
Closure called inside
@torch.no_grad(), preventing gradient computation
closure() is invoked while torch.no_grad() is active. Any loss.backward() call inside the closure will silently produce zero/no gradients. The standard PyTorch pattern (used in SGD, Adam, etc.) is to wrap the closure in with torch.enable_grad():.
| def step(self, closure=None): | |
| """Perform a single optimization step.""" | |
| loss = None | |
| if closure is not None: | |
| loss = closure() | |
| @torch.no_grad() | |
| def step(self, closure=None): | |
| """Perform a single optimization step.""" | |
| loss = None | |
| if closure is not None: | |
| with torch.enable_grad(): | |
| loss = closure() |
| scale_mode: str, | ||
| extra_scale_factor: float, | ||
| eps: float, | ||
| ) -> torch.Tensor: | ||
| global_shape = [grad.size(0), grad.size(1)] | ||
| global_shape[partition_dim] *= world_size |
There was a problem hiding this comment.
Reference
global_shape incorrectly scales an already-full tensor
_reference_orthogonalize receives the full matrix (shape full_shape) but then multiplies global_shape[partition_dim] by world_size a second time. For partition_dim=1 with world_size=2 and full_shape=(96, 128) this gives global_shape=[96, 256], so get_muon_scale_factor returns max(96,256)^0.5 = 16. The optimizer, operating on the shard (96, 64), correctly reconstructs global_shape=[96, 128] and computes max(96,128)^0.5 ≈ 11.3. This √2 discrepancy means the reference cannot correctly validate the optimizer's output.
The global_shape[partition_dim] *= world_size line should be removed since the input is already the full matrix.
| if mode == "unit_rms_norm": | ||
| return (size_out / size_in) ** 0.5 |
There was a problem hiding this comment.
unit_rms_norm mode can divide by zero when size_in == 0
(size_out / size_in) ** 0.5 raises ZeroDivisionError when size_in is 0. While the optimizer validates that the partition dimension is non-empty, it doesn't ensure the other dimension is non-zero. Consider adding a guard or documenting that both dimensions must be strictly positive.
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
skyw
left a comment
There was a problem hiding this comment.
I'd advice NOT to expose it in public API. Keeping it in test only if that is the purpose.
Having an optimizer with most code copied invites fragmentation.
Before this, all optimizer TE provides are more optimized fused version. I'd say a highly optimized Fused Muon with similar concept can be justified, but would need more consideration because it has more dependencies on other part of the training pipeline than elementwise optimizers.
| on tensor-parallel parameter shards. The local parameter shard must represent a | ||
| partition of a logical 2D matrix across the provided NCCL process group. | ||
|
|
||
| Args: |
There was a problem hiding this comment.
Q: Does TE use numpy style docstring instead of Google style?
|
|
||
| def __init__( | ||
| self, | ||
| params: Iterable[torch.nn.Parameter | dict], |
There was a problem hiding this comment.
Nit: The type here doesn't match PyTorch internal. Should be fine for the purpose of this class.
| scale_mode: MuonScaleT = "spectral", | ||
| extra_scale_factor: float = 1.0, | ||
| process_group: Optional[dist.ProcessGroup] = None, | ||
| partition_dim: int = 1, |
| raise ValueError(f"Invalid weight_decay value: {weight_decay}") | ||
| if num_ns_steps < 1: | ||
| raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") | ||
| if partition_dim not in (0, 1): |
There was a problem hiding this comment.
Q: Does this class intend to support non-distributed case? partition_dim would be -1 in TE in such case.
|
|
||
| if process_group is None: | ||
| if not dist.is_initialized(): | ||
| raise RuntimeError("MuonOptimizer requires torch.distributed to be initialized.") |
There was a problem hiding this comment.
Same question above regarding single GPU support.
| if process_group is None: | ||
| if not dist.is_initialized(): | ||
| raise RuntimeError("MuonOptimizer requires torch.distributed to be initialized.") | ||
| process_group = dist.group.WORLD |
There was a problem hiding this comment.
Suggestion: This silent behavior is dangerous. If user forgot to pass the correct TP group, wrong group will be used.
| global_shape[partition_dim] *= world_size | ||
|
|
||
| orth_grad = grad.clone() | ||
| transposed = partition_dim == 0 |
There was a problem hiding this comment.
Attn: This is from common Row and Column wise tensor parallelism in most LLM. It would be sub optimal for anything other than that. Add comment if the assumption is made.
The idea was to give something to users, who use TE, but not Megatron-LM. By fragmentation you mean that we want to encourage everyone to use Megatron-LM? Or that the optimizer being relatively thin thing on top of newton_schulz call, and the users should have no trouble creating it themselves? I don't think we gain anything by putting it into tests, since we already have tests for newton_schulz call. So we need to decide whether we want this PR, or should abandon it altogether. @cyanguwa |
Fragmentation means there will be different flavor of muon in emerging optimizer and TE, also a lot of copied code. TE can have stalled feature when emerging optimizer updates. Megatron-LM will always have its own version because there are implementation specific things need to be hooked together. For example, how QKV is implemetned, or fused swighlu. |
There was a problem hiding this comment.
Should we move newton_schulz.py to this directory? Also, how do we expect Megatron to call us for this functionality? Thanks.
There was a problem hiding this comment.
Should we move newton_schulz.py to this directory?
No, don't think so.
Should we move newton_schulz.py to this directory?
Megatron will call newton_shulz directly from their optimizers. This one is for other users.
There was a problem hiding this comment.
I'd prefer moving them into something like transformer_engine/pytorch/cusolver. But I suppose that is orthogonal to this PR.
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
|
|
@skyw Just following up on the discussion above - our purpose for this PR was two-fold. One was to provide an equivalent |
There was a problem hiding this comment.
I'd prefer moving them into something like transformer_engine/pytorch/cusolver. But I suppose that is orthogonal to this PR.
There was a problem hiding this comment.
We should make sure to include this in the QA script: https://github.com/NVIDIA/TransformerEngine/blob/main/qa/L1_pytorch_distributed_unittest/test.sh
| LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] | ||
|
|
||
|
|
||
| def _run_test(dtype: str, partition_dim: int, weight_decay_mode: str) -> None: |
There was a problem hiding this comment.
Each torchrun launch is somewhat expensive. Instead of launching a separate torchrun for each test case, it's better to launch a single torchrun instance and to perform multiple tests internally. See distributed/test_fusible_ops.py for an example.
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
Megatron is wrapped over emgering-optimizers with megatron specific details, like TP and how QKV are organized. The most optimizer logic is in emerging-optimizers. Could TE do the same? I understand introducing a new dependency may have concern, let me know. The biggest concern is actually large portion of duplicated code. What I would favor is having |
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
| partition_dim: Optional[int] = None, | ||
| tp_mode: Literal["duplicated", "distributed"] = "duplicated", | ||
| ) -> None: | ||
| """Compute tensor-parallel Newton-Schulz orthogonalization in-place. |
There was a problem hiding this comment.
Is there some requirement for the operation to be in-place? I see we are doing a few x.copy_()s in the code and was just wondering if it's necessary.
There was a problem hiding this comment.
Yes, it's intentional
| ctx : CusolverMpCtx | ||
| cuSolverMp context created for the tensor-parallel process group. | ||
| num_iterations : int, optional | ||
| Number of Newton-Schulz iterations. Default: 5. |
There was a problem hiding this comment.
Should this be optional, or mandatory? Shouldn't users always be specifying the exact number of iterations they want to run NS with? Otherwise, they might get successful runs but incorrect results silently?
|
|
||
| output_shards = [torch.empty_like(local_work) for _ in range(ctx.nranks)] | ||
| dist.all_gather(output_shards, local_work, group=ctx.group) | ||
| output = torch.cat(output_shards, dim=1) |
There was a problem hiding this comment.
I feel the functionality of _orthogonalize_replicated() is a bit odd. In the case of partition_dim is None, we probably don't need any chunking or all-gathering, given that tensor x is local, non-distributed. In the case of tp_mode == "duplicated", we seem to be all-gathering, chunking, and then all-gathering again - is this necessary? Looking at emerging-optimizers, it seems to be all lowering down to a simple newton_schulz() call?
There was a problem hiding this comment.
It's necessary for this cuSolverMp. cuSolverMp interprets each rank’s input as a column shard; the wrapper must shard replicated inputs before the call and gather the distributed result afterward. Emerging-Optimizers can avoid this only because its local PyTorch implementation is different.
| tp_mode=tp_mode, | ||
| ) | ||
| else: | ||
| newton_schulz(x_local, ctx, num_iterations, coefficients=coefficients) |
There was a problem hiding this comment.
Non-TP is a special case of TP, right? So we probably can just call newton_schulz_tp() here regardless of whether api == "tp"? We just need to set up partition_dim and tp_mode properly for the non-TP case?
There was a problem hiding this comment.
I think it makes sense to test both APIs, even if they end up running the same thing internally.
| return (5e-2, 5e-2) | ||
| if check == "orthogonality" and world_size == 1: | ||
| return (2e-2, 2e-2) | ||
| return (1e-2, 1e-2) |
There was a problem hiding this comment.
Are the tolerances set a bit high? Usually we have 1e-2 for BF16 and even smaller for FP32. Could you do a bit of research and tighten them up if necessary? A couple of reference points are:
https://github.com/NVIDIA-NeMo/Emerging-Optimizers/blob/main/tests/test_muon_utils.py
https://github.com/vcherepanov-nv/TransformerEngine/blob/3c4dfffb788f01009bf6741d25d4b69078e95d13/tests/pytorch/test_permutation.py#L192
There was a problem hiding this comment.
Are we committing to only supporting Newton-Shulz, and to only use it for Muon? It would be more general to put this code in a transformer_engine.pytorch.cusolvermp submodule.
There was a problem hiding this comment.
I think in principle that's a good idea, but if there's realistically no other use case of it in TE, it might be better to simplify the file structure?
There was a problem hiding this comment.
Could you check if there are other optimizers in Emerging-Optimizers that could be using cusolvermp?
| tp_mode=tp_mode, | ||
| ) | ||
| else: | ||
| newton_schulz(x_local, ctx, num_iterations, coefficients=coefficients) |
There was a problem hiding this comment.
I think it makes sense to test both APIs, even if they end up running the same thing internally.
# Conflicts: # tests/pytorch/distributed/run_newton_schulz.py # transformer_engine/pytorch/optimizers/newton_schulz.py
Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>
for more information, see https://pre-commit.ci
cyanguwa
left a comment
There was a problem hiding this comment.
Looks like the CI is clean as well. Thanks!
* Add distributed Muon optimizer Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Muon closure and reference test Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer distributed API handling Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer docs and params typing Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tensor-parallel Newton-Schulz wrapper Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Move Newton-Schulz wrapper into optimizers Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Use tensor-parallel Newton-Schulz in Muon Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Limit Muon branch to Newton-Schulz TP Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Run Newton-Schulz distributed cases in one launch Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Add single-GPU Newton-Schulz coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Newton-Schulz compatibility and replicated coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> (cherry picked from commit 5df5a0d)
Description
Add a newton_schulz_tp utility function on top of the existing newton_schulz. To be used in Muon optimizer step.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: