Microsoft Dion: orthonormal optimizers that assume DTensor parallelism

What an orthonormal optimizer changes about the update rule

The Introduction section frames the problem in the usual terms: optimization algorithms are essential to training neural networks, converting gradients into model weight updates to minimize loss. For many years, the method of choice has been Adam and AdamW. However, recent work has shown that orthonormal optimizers can significantly accelerate model convergence.

So the category is not a new loss function or a new architecture. It is a different way of turning a gradient into an update, one that constrains the update to lie in an orthonormal basis rather than applying a per-coordinate scaled step.

The repository ships four named optimizers, and the distinctions between them matter more than the shared framing. Muon is the original, with a blog post by Keller Jordan and a paper. Dion2 and Dion3 are the current Microsoft contributions, linked to an arXiv preprint, and Dion itself is called a legacy optimizer with the README explicitly recommending Dion2 or Dion3 instead. NorMuon is a fourth variant with its own arXiv entry.

The README credits the practical effectiveness of the category to work by Jeremy Bernstein and Laker Newhouse, and says it was first demonstrated by Muon in the NanoGPT speedrun, and has since been validated at scale by models such as Kimi K2.

That last point is the one to hold on to. The claim being made is not that an orthonormal optimizer helps on a small model. It is that this approach has been used at the scale of a frontier model, which is a much stronger statement about the method than a benchmark on a small setup would be.

The primary language is Python, the licence is MIT with a LICENSE file and a NOTICE.md at the root, the default branch is main, and the last push was on 2026-09-30. There are no GitHub releases, so the package version lives in setup.py rather than in a tag.

The PyTorch requirement is DTensor, and it rules out other training APIs

The Requirements section is one sentence and it is a hard constraint rather than a preference.

This code is written for modern PyTorch, version 2.7 or newer, using DTensor-based parallelism. This includes FSDP2 with fully_shard and tensor parallelism with parallelize_module. Support for other distributed training APIs is not implemented.

That sentence defines the addressable audience. If your training stack already uses FSDP2 or tensor parallelism through DTensor, Dion fits into it. If your stack uses a different distributed API, the library does not work with it, and that is a design decision rather than an oversight.

The reason is visible in the repository layout. The dion/ directory holds the optimizer implementations, and the design assumes that a parameter is a DTensor sharded across a mesh. An orthonormal optimizer has to orthogonalise across the full matrix of weights, and with sharding that means the update computation has to be aware of how the tensor is distributed. A library written for DTensor gets that for free; a library written for a different API would need its own distribution-aware path.

The README also warns that there are major differences compared to PyTorch built-in optimizers such as Adam and AdamW, and asks readers to go through the README carefully for detailed instructions. An optimizer that replaces AdamW in a training script is not a drop-in swap, because the parameter groups are built differently, and the README has a whole section on it.

That is where building parameter groups comes in. The optimizers operate on parameter groups you construct rather than on whatever an optimizer factory hands you, with per-head Newton-Schulz for attention projections, and the README documents the parameter group construction as its own section with example code.

The version floor of PyTorch 2.7 is also not arbitrary. DTensor and FSDP2 both arrived in the 2.x series, and the fully_shard API in particular was still settling in the releases before that.

Installing from git, with the Cutlass kernels kept behind an extra

The installation section is short, and the detail that matters is the extras split.

The base install is one command.

pip install git+https://github.com/microsoft/dion.git

Then the import line, which tells you the whole public surface.

from dion import Dion2, Dion3, Muon, NorMuon, Dion

What follows is the part worth reading twice. The optional Gram Newton-Schulz orthogonalization kernels, enabled with use_gram_newton_schulz=True, are not pulled in by the base install. You add them with the gram-newton-schulz extra, or with pip install -e ".[gram-newton-schulz]" from a clone.

So the kernel choice is not a runtime flag alone, it is a different install. That is a deliberate separation: the optional path needs a specific CUDA toolchain, and forcing everyone to have it would make the base install heavier and more fragile.

The extra pins nvidia-cutlass-dsl==4.5.2, matching the cutlass version used by current Flash-Attention-4 and Blackwell stacks, and the README notes that the earlier 4.4.2 conflict no longer applies. That is a small but revealing detail: the pinned version tracks what current GPU stacks use, so on a machine with a different Cutlass version in the environment you will find out at install time rather than at training time.

For development and training, the repository has separate extras. The train extra installs the dependencies for the training code and stays free of the Gram Newton-Schulz kernels, with pip install -e ".[train,gns]" for training with the kernels. The dev extra skips the GNS-specific test cases in a plain install.

Looking at setup.py explains how this is wired. There are four requirements files, requirements_dion.txt for the standalone optimizer, requirements_gns.txt for the optional kernels, requirements_dev.txt and requirements_train.txt, and the extras are assembled by concatenating the base list with the relevant one. The version is hardcoded in setup.py as 0.1.0, with a comment noting that versions below 1.0 are considered pre-release, and a commented-out block showing how to read the version from a file instead if you use a tool such as bump2version.

Sharding combinations are spelled out, and their product must equal the world size

The distributed training section is the most useful part of the README for anyone actually running this, because it gives concrete invocations for each parallelism strategy and states the constraint that ties them together.

Distributed Data Parallel training on four GPUs is the baseline.

torchrun --standalone --nproc_per_node=4 train.py --config configs/dion2_160m.yaml

FSDP is enabled by giving the size of the FSDP group with --fs_size.

torchrun --standalone --nproc_per_node=4 train.py \
  --config configs/dion2_160m.yaml \
  --fs_size 4

Hybrid Sharded Data Parallel sets --fs_size smaller than the GPU count and gives the data parallel dimension with --dp_size.

torchrun --standalone --nproc_per_node=4 train.py \
  --config configs/dion2_160m.yaml \
  --fs_size 2 \
  --dp_size 2

That produces two FSDP groups each spanning two GPUs, with two-way data parallelism across the groups, four GPUs in total.

The rule that ties them together is stated plainly: the product dp_size times fs_size must equal world_size, and any unspecified dimension defaults to 1.

Tensor parallelism is the third axis, and it carries an important restriction. Currently, only Dion, the legacy implementation, supports Tensor Parallelism. A two by two by two configuration across eight GPUs combines all three strategies.

torchrun --standalone --nproc_per_node=8 train.py \
  --config configs/dion_160m.yaml \
  --dp_size 2 \
  --fs_size 2 \
  --tp_size 2

Here the product dp_size times fs_size times tp_size must equal world_size.

So the practical summary is: Dion2 and Dion3 work with DDP, FSDP and HSDP, while the legacy Dion is the one that also works with tensor parallelism. If you need TP, you are on the legacy optimizer, and the README recommends against that for new work.

Dion3 is not a retuning of Dion2, and the README says so

There is a paragraph in the DDP section that deserves more attention than it usually gets, because it is the kind of thing that quietly breaks a training run.

To train with Dion3 instead, swap in configs/dion3_160m.yaml. Its hyperparameters are carried over from the Dion2 config as a starting point rather than a translation, because Dion3's per-neuron normalization redistributes the update across rows, and mu feeds Dion2's error-feedback decay but Dion3's momentum, so you are told to read the comments in the file before tuning.

Three things are being said there. First, the Dion3 config starts from Dion2's values because there is no principled translation available. Second, the reason a translation is not available is that the update itself changes shape: Dion3 normalises per neuron and redistributes the update across rows, so the same mu means a different thing. Third, the momentum parameter mu is overloaded between the two optimizers, doing error-feedback decay in Dion2 and momentum in Dion3, which is exactly the kind of collision that produces a training run that looks plausible and converges badly.

The README then says the FSDP and HSDP examples accept it as well, and that tensor parallelism does not, which is consistent with the earlier note that only the legacy Dion supports TP.

This is a good example of documentation that tells you the cost of the feature rather than only the capability. The cost is that swapping optimizers is not a config change; it is a re-tuning exercise with a specific trap in it.

The repository layout supports this workflow. The configs/ directory holds the yaml files, train.py is the training script, data/ has the dataset preparation including cached_fineweb10B.py, models/ holds the model definitions, and there are tests/, a benchmark/ directory and separate requirements files for the dev, train and GNS paths. The sample data preparation step is a single command, python data/cached_fineweb10B.py 30, which downloads the pretokenized FineWeb dataset.

Where Dion is the wrong tool

The first wrong-tool case is single-GPU training or a project on a distributed API other than DTensor. Support for other distributed training APIs is not implemented, and that is not a small gap: an orthonormal optimizer has to be distribution-aware about how weights are sharded, and the library gets that from DTensor specifically. A team on a legacy DistributedDataParallel or a non-PyTorch training stack is out of options here.

The second case is a project that wants the AdamW API. The README says there are major differences compared to PyTorch built-in optimizers, and building parameter groups is a separate documented step with per-head Newton-Schulz for attention projections. An optimizer that swaps in for AdamW usually means rewriting the parameter grouping, which for a large model with mixed tensor shapes is not a small change. If the goal is a small experiment on an existing AdamW script, the friction is real.

The third case is a team that needs a stable, versioned release. There are no GitHub releases, the package version in setup.py is 0.1.0, and the code comments note that versions below 1.0 are considered pre-release. The install is pip install from a git URL rather than from PyPI, which means you are installing whatever the branch head is at the moment you run it. There is a CHANGELOG.md in the repository, but a changelog without tags is a weaker guarantee than a version you can pin.

There is also a hardware coupling to check. The optional Gram Newton-Schulz kernels pin nvidia-cutlass-dsl==4.5.2 to match current Flash-Attention-4 and Blackwell stacks. That is deliberate, and it means the accelerated path is tied to a specific Cutlass version, so on a different CUDA or GPU generation the install either needs the matching extra or has to fall back to the non-GNS path.

On licence, MIT with a LICENSE file and a NOTICE.md at the root, which is the permissive grant with no copyleft obligation on a user. The licence does not, however, cover the models you train or the datasets you train on, and those carry their own terms.

What the repository gives you is a set of optimizer implementations that are built for the current PyTorch distributed stack rather than bolted on afterwards. If you are already running FSDP2 or tensor parallelism through DTensor and you want to try an orthonormal optimizer, the code is written for exactly that, with the sharding combinations documented and the Dion2 versus Dion3 differences called out. If you are not, the friction is the DTensor requirement, and it is a hard one.

Editorial conclusion

Dion suits a PyTorch team training on multiple GPUs that wants to try an orthonormal optimizer such as Muon or Dion3 in an existing distributed setup, because the implementations are written against the same DTensor parallelism the rest of a modern PyTorch stack uses and the README documents the sharding combinations in detail. It does not suit a single-GPU project, and it does not suit a stack built on a distributed API other than DTensor, since support for those is explicitly not implemented. Before adopting, read the Dion2 versus Dion3 configuration notes, because the README warns that Dion3's per-neuron normalization changes how the update is distributed and that the hyperparameters are carried over rather than translated.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. microsoft/dion on GitHub
4. README